feat: initial HSAP platform
Huaxu Sentinel Active Safety Platform with embedded algorithm code, Docker Compose setup, and vendored dataset scaffolds for clone-and-run. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
50
algorithms/dms_yolo/code/examples/README.md
Normal file
50
algorithms/dms_yolo/code/examples/README.md
Normal file
@@ -0,0 +1,50 @@
|
||||
<a href="https://www.ultralytics.com/" target="_blank"><img src="https://raw.githubusercontent.com/ultralytics/assets/main/logo/Ultralytics_Logotype_Original.svg" width="320" alt="Ultralytics logo"></a>
|
||||
|
||||
# Ultralytics YOLO Examples
|
||||
|
||||
> [!WARNING]
|
||||
> The examples in this directory are community-contributed and showcase creative ways to use Ultralytics YOLO models. While we truly appreciate these contributions, they may not always reflect the latest best practices or receive regular updates. To help streamline our codebase and focus our resources on maintaining comprehensive, up-to-date official documentation and guides, we plan to retire these examples in **Ultralytics v8.4.0**.
|
||||
|
||||
Welcome to the Ultralytics examples directory! This collection showcases practical applications and detailed walkthroughs for integrating [Ultralytics YOLO models](https://docs.ultralytics.com/models/) into various real-world projects. Explore Python scripts and Jupyter notebooks designed to help you leverage the power of models like [Ultralytics YOLO26](https://docs.ultralytics.com/models/yolo26/) for tasks like [object detection](https://docs.ultralytics.com/tasks/detect/), [instance segmentation](https://docs.ultralytics.com/tasks/segment/), [pose estimation](https://docs.ultralytics.com/tasks/pose/), and more.
|
||||
|
||||
Whether you're deploying models on [edge devices](https://www.ultralytics.com/glossary/edge-ai) using formats like [ONNX](https://docs.ultralytics.com/integrations/onnx/) with [ONNX Runtime](https://onnxruntime.ai/), optimizing with [TensorRT](https://docs.ultralytics.com/integrations/tensorrt/) on NVIDIA Jetson, using [OpenVINO](https://docs.ultralytics.com/integrations/openvino/) for Intel hardware, or integrating with frameworks like [OpenCV](https://opencv.org/), these examples provide valuable insights and code snippets. Find inspiration for your next [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) project and see how others are using Ultralytics YOLO to build innovative [AI solutions](https://www.ultralytics.com/solutions) on platforms ranging from C++ and C# to Python and Rust.
|
||||
|
||||
## 💡 Example Applications
|
||||
|
||||
Browse through the community-contributed examples below. These projects demonstrate various use cases and deployment strategies for Ultralytics YOLO models across different platforms and programming languages.
|
||||
|
||||
| Title | Format | Contributor |
|
||||
| ----------------------------------------------------------------------------------------------------------------------------------------- | -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| [YOLO ONNX Detection Inference with C++](./YOLOv8-CPP-Inference) | C++/ONNX | [Justas Bartnykas](https://github.com/JustasBart) |
|
||||
| [YOLO OpenCV ONNX Detection Python](./YOLOv8-OpenCV-ONNX-Python) | OpenCV/Python/ONNX | [Farid Inawan](https://github.com/frdteknikelektro) |
|
||||
| [YOLO C# ONNX-Runtime](https://github.com/dme-compunet/YoloSharp) | .NET/ONNXRuntime | [Compunet](https://github.com/dme-compunet) |
|
||||
| [YOLO .Net ONNX Detection C#](https://www.nuget.org/packages/Yolov8.Net) | C# .Net | [Samuel Stainback](https://github.com/sstainba) |
|
||||
| [YOLOv8 on NVIDIA Jetson (TensorRT and DeepStream)](https://wiki.seeedstudio.com/YOLOv8-DeepStream-TRT-Jetson/) | Python/TensorRT/DeepStream | [Lakshantha](https://github.com/lakshanthad) ([See also DeepStream Guide](https://docs.ultralytics.com/guides/deepstream-nvidia-jetson/)) |
|
||||
| [YOLOv8 ONNXRuntime Python](./YOLOv8-ONNXRuntime) | Python/ONNXRuntime | [Semih Demirel](https://github.com/semihhdemirel) |
|
||||
| [RTDETR ONNXRuntime Python](./RTDETR-ONNXRuntime-Python) | Python/ONNXRuntime | [Semih Demirel](https://github.com/semihhdemirel) |
|
||||
| [YOLOv8 ONNXRuntime CPP](./YOLOv8-ONNXRuntime-CPP) | C++/ONNXRuntime | [DennisJcy](https://github.com/DennisJcy), [Onuralp Sezer](https://github.com/onuralpszr) |
|
||||
| [RTDETR ONNXRuntime C#](https://github.com/Kayzwer/yolo-cs/blob/master/RTDETR.cs) | C#/ONNX | [Kayzwer](https://github.com/Kayzwer) |
|
||||
| [YOLOv8 SAHI Video Inference](https://github.com/RizwanMunawar/ultralytics/blob/main/examples/YOLOv8-SAHI-Inference-Video/yolov8_sahi.py) | Python/SAHI | [Muhammad Rizwan Munawar](https://github.com/RizwanMunawar) ([See also SAHI Guide](https://docs.ultralytics.com/guides/sahi-tiled-inference/)) |
|
||||
| [YOLOv8 Region Counter](https://github.com/RizwanMunawar/ultralytics/blob/main/examples/YOLOv8-Region-Counter/yolov8_region_counter.py) | Python | [Muhammad Rizwan Munawar](https://github.com/RizwanMunawar) ([See also Region Counting Guide](https://docs.ultralytics.com/guides/region-counting/)) |
|
||||
| [YOLOv8 Segmentation ONNXRuntime Python](./YOLOv8-Segmentation-ONNXRuntime-Python) | Python/ONNXRuntime | [jamjamjon](https://github.com/jamjamjon) |
|
||||
| [YOLOv8 LibTorch CPP](./YOLOv8-LibTorch-CPP-Inference) | C++/LibTorch | [Myyura](https://github.com/Myyura) ([See also TorchScript Export](https://docs.ultralytics.com/integrations/torchscript/)) |
|
||||
| [YOLOv8 OpenCV INT8 TFLite Python](./YOLOv8-TFLite-Python) | Python/TFLite | [Wamiq Raza](https://github.com/wamiqraza) ([See also TFLite Export](https://docs.ultralytics.com/integrations/tflite/)) |
|
||||
| [YOLOv8 All Tasks ONNXRuntime Rust](./YOLOv8-ONNXRuntime-Rust) | Rust/ONNXRuntime | [jamjamjon](https://github.com/jamjamjon) |
|
||||
| [YOLOv8 OpenVINO CPP](./YOLOv8-OpenVINO-CPP-Inference) | C++/OpenVINO | [Erlangga Yudi Pradana](https://github.com/rlggyp) ([See also OpenVINO Export](https://docs.ultralytics.com/integrations/openvino/)) |
|
||||
| [YOLOv8 MNN CPP](./YOLOv8-MNN-CPP) | C++/MNN | [Khoi VN](https://github.com/vnk8071) |
|
||||
| [YOLOv5-YOLO11 ONNXRuntime Rust](./YOLO-Series-ONNXRuntime-Rust) | Rust/ONNXRuntime | [jamjamjon](https://github.com/jamjamjon) |
|
||||
|
||||
## 🤝 How to Contribute
|
||||
|
||||
We actively encourage contributions from our vibrant community! Sharing your examples, applications, and guides helps others learn and build amazing things with [Ultralytics](https://www.ultralytics.com/). If you have a project you'd like to share, please follow these steps:
|
||||
|
||||
1. **Fork the Repository:** Start by forking the main [Ultralytics repository](https://github.com/ultralytics/ultralytics) on [GitHub](https://github.com/).
|
||||
2. **Create Your Example:** Add your project folder within the `examples/` directory of your forked repository.
|
||||
3. **Prepare Your Submission:** Ensure your project meets the following criteria:
|
||||
- It utilizes the `ultralytics` pip package.
|
||||
- Includes a `README.md` file with clear, step-by-step instructions for setup and execution. Explain the purpose of the example and any prerequisites.
|
||||
- Avoid committing large files or extensive dependencies. If necessary, provide instructions for users to download them separately (e.g., using `ultralytics.utils.downloads.safe_download()`).
|
||||
- As a contributor, be prepared to offer support and address [issues](https://github.com/ultralytics/ultralytics/issues) related to your example.
|
||||
4. **Submit a Pull Request:** Create a [pull request (PR)](https://github.com/ultralytics/ultralytics/pulls) targeting the `main` branch of the official Ultralytics repository. Use the title prefix `[Example]` (e.g., `[Example] Add YOLOv8 Pose Estimation on Raspberry Pi`).
|
||||
|
||||
For more comprehensive guidelines on contributing code, documentation, or examples, please refer to our [Contributing Guide](https://docs.ultralytics.com/help/contributing/). We appreciate your efforts to enhance the Ultralytics ecosystem! If you have questions, feel free to open an issue or PR, and the team will be happy to assist. Check out the [Ultralytics Blog](https://www.ultralytics.com/blog) for more insights and updates, and explore [Ultralytics Platform](https://platform.ultralytics.com/ultralytics/yolo26) for streamlined model training and deployment.
|
||||
@@ -0,0 +1,46 @@
|
||||
# RT-DETR Object Detection with ONNX Runtime
|
||||
|
||||
This project demonstrates how to run Ultralytics [RT-DETR models](https://docs.ultralytics.com/models/rtdetr/) using the [ONNX Runtime](https://onnxruntime.ai/) inference engine in [Python](https://www.python.org/). It provides a straightforward example for performing [object detection](https://docs.ultralytics.com/tasks/detect/) with RT-DETR models that have been exported to the [ONNX format](https://onnx.ai/), a standard for representing [machine learning models](https://www.ultralytics.com/glossary/machine-learning-ml). RT-DETR, or Real-Time DEtection TRansformer, offers efficient and accurate object detection capabilities, detailed further in the [RT-DETR research paper](https://arxiv.org/abs/2304.08069).
|
||||
|
||||
## ⚙️ Installation
|
||||
|
||||
To get started, you'll need to install the necessary dependencies. Follow the steps below.
|
||||
|
||||
### Installing Required Dependencies
|
||||
|
||||
Install the core requirements using [pip](https://pip.pypa.io/en/stable/) and the provided `requirements.txt` file. This installs the standard **`onnxruntime`** package (CPU-based inference). See the [ONNX Runtime Execution Providers documentation](https://onnxruntime.ai/docs/execution-providers/) for more information on available execution options.
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### Installing `onnxruntime-gpu` (Optional)
|
||||
|
||||
For accelerated inference using an NVIDIA GPU, install the **`onnxruntime-gpu`** package. Ensure you have the correct [NVIDIA drivers](https://www.nvidia.com/Download/index.aspx) and [CUDA toolkit](https://developer.nvidia.com/cuda-toolkit) installed first. Consult the official [ONNX Runtime GPU documentation](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html) for detailed compatibility information and setup instructions.
|
||||
|
||||
```bash
|
||||
pip install onnxruntime-gpu
|
||||
```
|
||||
|
||||
## 🚀 Usage
|
||||
|
||||
Once the dependencies are installed, you can run inference using the `main.py` script.
|
||||
|
||||
Execute the script from your terminal, specifying the path to your ONNX model, the input image, and optional confidence and IoU thresholds:
|
||||
|
||||
```bash
|
||||
python main.py --model rtdetr-l.onnx --img image.jpg --conf-thres 0.5 --iou-thres 0.5
|
||||
```
|
||||
|
||||
**Arguments:**
|
||||
|
||||
- `--model`: Path to the RT-DETR [ONNX model file](https://docs.ultralytics.com/modes/export/) (e.g., `rtdetr-l.onnx`). You can easily [export Ultralytics models](https://docs.ultralytics.com/modes/export/) to ONNX format. Find more models on the [Ultralytics Models](https://docs.ultralytics.com/models/) page.
|
||||
- `--img`: Path to the input image file (e.g., `image.jpg`).
|
||||
- `--conf-thres`: Confidence threshold for filtering detections. Only detections with a score higher than this value will be kept. Learn more about thresholds in our guide on [YOLO performance metrics](https://docs.ultralytics.com/guides/yolo-performance-metrics/).
|
||||
- `--iou-thres`: [Intersection over Union (IoU)](https://www.ultralytics.com/glossary/intersection-over-union-iou) threshold used for [Non-Maximum Suppression (NMS)](https://www.ultralytics.com/glossary/non-maximum-suppression-nms) to remove redundant [bounding boxes](https://www.ultralytics.com/glossary/bounding-box).
|
||||
|
||||
Adjust the `--conf-thres` and `--iou-thres` values based on your specific requirements for detection sensitivity and overlap removal.
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions to enhance this example are welcome! Whether it's fixing bugs, adding new features, improving documentation, or suggesting optimizations, your input is valuable. Please refer to the Ultralytics [Contribution Guide](https://docs.ultralytics.com/help/contributing/) for detailed information on how to get started. You can also explore general guides on [contributing to open source projects](https://opensource.guide/how-to-contribute/). Thank you for helping improve the [Ultralytics](https://www.ultralytics.com/) ecosystem and its resources available on [GitHub](https://github.com/ultralytics/ultralytics)!
|
||||
@@ -0,0 +1,298 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import requests
|
||||
import yaml
|
||||
|
||||
|
||||
def download_file(url: str, local_path: str) -> str:
|
||||
"""Download a file from a URL to a local path.
|
||||
|
||||
Args:
|
||||
url (str): URL of the file to download.
|
||||
local_path (str): Local path where the file will be saved.
|
||||
"""
|
||||
# Check if the local path already exists
|
||||
if os.path.exists(local_path):
|
||||
print(f"File already exists at {local_path}. Skipping download.")
|
||||
return local_path
|
||||
# Download the file from the URL
|
||||
print(f"Downloading {url} to {local_path}...")
|
||||
response = requests.get(url, stream=True, timeout=30)
|
||||
response.raise_for_status()
|
||||
with open(local_path, "wb") as f:
|
||||
for chunk in response.iter_content(chunk_size=1024 * 1024):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
return local_path
|
||||
|
||||
|
||||
class RTDETR:
|
||||
"""RT-DETR (Real-Time Detection Transformer) object detection model for ONNX inference and visualization.
|
||||
|
||||
This class implements the RT-DETR model for object detection tasks, supporting ONNX model inference and
|
||||
visualization of detection results with bounding boxes and class labels.
|
||||
|
||||
Attributes:
|
||||
model_path (str): Path to the ONNX model file.
|
||||
img_path (str): Path to the input image.
|
||||
conf_thres (float): Confidence threshold for filtering detections.
|
||||
iou_thres (float): IoU threshold for non-maximum suppression.
|
||||
session (ort.InferenceSession): ONNX runtime inference session.
|
||||
model_input (list): Model input metadata.
|
||||
input_width (int): Width dimension required by the model.
|
||||
input_height (int): Height dimension required by the model.
|
||||
classes (list[str]): List of class names from COCO dataset.
|
||||
color_palette (np.ndarray): Random color palette for visualization.
|
||||
img (np.ndarray): Loaded input image.
|
||||
img_height (int): Height of the input image.
|
||||
img_width (int): Width of the input image.
|
||||
|
||||
Methods:
|
||||
draw_detections: Draw bounding boxes and labels on the input image.
|
||||
preprocess: Preprocess the input image for model inference.
|
||||
bbox_cxcywh_to_xyxy: Convert bounding boxes from center format to corner format.
|
||||
postprocess: Postprocess model output to extract and visualize detections.
|
||||
main: Execute the complete object detection pipeline.
|
||||
|
||||
Examples:
|
||||
Initialize RT-DETR detector and run inference
|
||||
>>> detector = RTDETR("rtdetr-l.onnx", "image.jpg", conf_thres=0.5)
|
||||
>>> output_image = detector.main()
|
||||
>>> cv2.imshow("Detections", output_image)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
img_path: str,
|
||||
conf_thres: float = 0.5,
|
||||
iou_thres: float = 0.5,
|
||||
class_names: str | None = None,
|
||||
):
|
||||
"""Initialize the RT-DETR object detection model.
|
||||
|
||||
Args:
|
||||
model_path (str): Path to the ONNX model file.
|
||||
img_path (str): Path to the input image.
|
||||
conf_thres (float, optional): Confidence threshold for filtering detections.
|
||||
iou_thres (float, optional): IoU threshold for non-maximum suppression.
|
||||
class_names (Optional[str], optional): Path to a YAML file containing class names. If None, uses COCO
|
||||
dataset classes.
|
||||
"""
|
||||
self.model_path = model_path
|
||||
self.img_path = img_path
|
||||
self.conf_thres = conf_thres
|
||||
self.iou_thres = iou_thres
|
||||
self.classes = class_names
|
||||
|
||||
# Set up the ONNX runtime session with available execution providers
|
||||
available = ort.get_available_providers()
|
||||
providers = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in available]
|
||||
self.session = ort.InferenceSession(model_path, providers=providers or available)
|
||||
|
||||
self.model_input = self.session.get_inputs()
|
||||
self.input_width = self.model_input[0].shape[2]
|
||||
self.input_height = self.model_input[0].shape[3]
|
||||
|
||||
if self.classes is None:
|
||||
# Load class names from the COCO dataset YAML file
|
||||
self.classes = download_file(
|
||||
"https://raw.githubusercontent.com/ultralytics/ultralytics/main/ultralytics/cfg/datasets/coco8.yaml",
|
||||
"coco8.yaml",
|
||||
)
|
||||
|
||||
# Parse the YAML file to get class names
|
||||
with open(self.classes) as f:
|
||||
class_data = yaml.safe_load(f)
|
||||
self.classes = list(class_data["names"].values())
|
||||
|
||||
# Ensure the classes are a list
|
||||
if not isinstance(self.classes, list):
|
||||
raise ValueError("Classes should be a list of class names.")
|
||||
|
||||
# Generate a color palette for drawing bounding boxes
|
||||
self.color_palette: np.ndarray = np.random.uniform(0, 255, size=(len(self.classes), 3))
|
||||
|
||||
def draw_detections(self, box: np.ndarray, score: float, class_id: int) -> None:
|
||||
"""Draw bounding box and label on the input image for a detected object."""
|
||||
# Extract the coordinates of the bounding box
|
||||
x1, y1, x2, y2 = box
|
||||
|
||||
# Retrieve the color for the class ID
|
||||
color = self.color_palette[class_id]
|
||||
|
||||
# Draw the bounding box on the image
|
||||
cv2.rectangle(self.img, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
|
||||
|
||||
# Create the label text with class name and score
|
||||
label = f"{self.classes[class_id]}: {score:.2f}"
|
||||
|
||||
# Calculate the dimensions of the label text
|
||||
(label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||||
|
||||
# Calculate the position of the label text
|
||||
label_x = x1
|
||||
label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10
|
||||
|
||||
# Draw a filled rectangle as the background for the label text
|
||||
cv2.rectangle(
|
||||
self.img,
|
||||
(int(label_x), int(label_y - label_height)),
|
||||
(int(label_x + label_width), int(label_y + label_height)),
|
||||
color,
|
||||
cv2.FILLED,
|
||||
)
|
||||
|
||||
# Draw the label text on the image
|
||||
cv2.putText(
|
||||
self.img, label, (int(label_x), int(label_y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA
|
||||
)
|
||||
|
||||
def preprocess(self) -> np.ndarray:
|
||||
"""Preprocess the input image for model inference.
|
||||
|
||||
Loads the image, converts color space from BGR to RGB, resizes to model input dimensions, and normalizes pixel
|
||||
values to [0, 1] range.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Preprocessed image data with shape (1, 3, H, W) ready for inference.
|
||||
"""
|
||||
# Read the input image using OpenCV
|
||||
self.img = cv2.imread(self.img_path)
|
||||
if self.img is None:
|
||||
raise FileNotFoundError(f"Image not found or unreadable: '{self.img_path}'")
|
||||
|
||||
# Get the height and width of the input image
|
||||
self.img_height, self.img_width = self.img.shape[:2]
|
||||
|
||||
# Convert the image color space from BGR to RGB
|
||||
img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# Resize the image to match the input shape
|
||||
img = cv2.resize(img, (self.input_width, self.input_height))
|
||||
|
||||
# Normalize the image data by dividing it by 255.0
|
||||
image_data = np.array(img) / 255.0
|
||||
|
||||
# Transpose the image to have the channel dimension as the first dimension
|
||||
image_data = np.transpose(image_data, (2, 0, 1)) # Channel first
|
||||
|
||||
# Expand the dimensions of the image data to match the expected input shape
|
||||
image_data = image_data[None].astype(np.float32)
|
||||
|
||||
return image_data
|
||||
|
||||
def bbox_cxcywh_to_xyxy(self, boxes: np.ndarray) -> np.ndarray:
|
||||
"""Convert bounding boxes from center format to corner format.
|
||||
|
||||
Args:
|
||||
boxes (np.ndarray): Array of shape (N, 4) where each row represents a bounding box in (center_x, center_y,
|
||||
width, height) format.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Array of shape (N, 4) with bounding boxes in (x_min, y_min, x_max, y_max) format.
|
||||
"""
|
||||
# Calculate half width and half height of the bounding boxes
|
||||
half_width = boxes[:, 2] / 2
|
||||
half_height = boxes[:, 3] / 2
|
||||
|
||||
# Calculate the coordinates of the bounding boxes
|
||||
x_min = boxes[:, 0] - half_width
|
||||
y_min = boxes[:, 1] - half_height
|
||||
x_max = boxes[:, 0] + half_width
|
||||
y_max = boxes[:, 1] + half_height
|
||||
|
||||
# Return the bounding boxes in (x_min, y_min, x_max, y_max) format
|
||||
return np.column_stack((x_min, y_min, x_max, y_max))
|
||||
|
||||
def postprocess(self, model_output: list[np.ndarray]) -> np.ndarray:
|
||||
"""Postprocess model output to extract and visualize detections.
|
||||
|
||||
Applies confidence thresholding, converts bounding box format, scales coordinates to original image dimensions,
|
||||
and draws detection annotations.
|
||||
|
||||
Args:
|
||||
model_output (list[np.ndarray]): Output tensors from the model inference.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Annotated image with detection bounding boxes and labels.
|
||||
"""
|
||||
# Squeeze the model output to remove unnecessary dimensions
|
||||
outputs = np.squeeze(model_output[0])
|
||||
|
||||
# Extract bounding boxes and scores from the model output
|
||||
boxes = outputs[:, :4]
|
||||
scores = outputs[:, 4:]
|
||||
|
||||
# Get the class labels and scores for each detection
|
||||
labels = np.argmax(scores, axis=1)
|
||||
scores = np.max(scores, axis=1)
|
||||
|
||||
# Apply confidence threshold to filter out low-confidence detections
|
||||
mask = scores > self.conf_thres
|
||||
boxes, scores, labels = boxes[mask], scores[mask], labels[mask]
|
||||
|
||||
# Convert bounding boxes to (x_min, y_min, x_max, y_max) format
|
||||
boxes = self.bbox_cxcywh_to_xyxy(boxes)
|
||||
|
||||
# Scale bounding boxes to match the original image dimensions
|
||||
boxes[:, 0::2] *= self.img_width
|
||||
boxes[:, 1::2] *= self.img_height
|
||||
|
||||
# Apply non-maximum suppression (optional for RT-DETR, but useful for filtering overlaps)
|
||||
xywh_boxes = [[float(b[0]), float(b[1]), float(b[2] - b[0]), float(b[3] - b[1])] for b in boxes]
|
||||
indices = cv2.dnn.NMSBoxes(xywh_boxes, scores.tolist(), self.conf_thres, self.iou_thres)
|
||||
indices = indices.flatten().tolist() if len(indices) else []
|
||||
|
||||
# Draw detections on the image
|
||||
for i in indices:
|
||||
self.draw_detections(boxes[i], float(scores[i]), int(labels[i]))
|
||||
|
||||
return self.img
|
||||
|
||||
def main(self) -> np.ndarray:
|
||||
"""Execute the complete object detection pipeline on the input image.
|
||||
|
||||
Performs preprocessing, ONNX model inference, and postprocessing to generate annotated detection results.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Output image with detection annotations including bounding boxes and class labels.
|
||||
"""
|
||||
# Preprocess the image for model input
|
||||
image_data = self.preprocess()
|
||||
|
||||
# Run the model inference
|
||||
model_output = self.session.run(None, {self.model_input[0].name: image_data})
|
||||
|
||||
# Process and return the model output
|
||||
return self.postprocess(model_output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Set up argument parser for command-line arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", type=str, default="rtdetr-l.onnx", help="Path to the ONNX model file.")
|
||||
parser.add_argument("--img", type=str, default="bus.jpg", help="Path to the input image.")
|
||||
parser.add_argument("--conf-thres", type=float, default=0.5, help="Confidence threshold for object detection.")
|
||||
parser.add_argument("--iou-thres", type=float, default=0.5, help="IoU threshold for non-maximum suppression.")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Create the detector instance with specified parameters
|
||||
detection = RTDETR(args.model, args.img, args.conf_thres, args.iou_thres)
|
||||
|
||||
# Perform detection and get the output image
|
||||
output_image = detection.main()
|
||||
|
||||
# Display the annotated output image
|
||||
cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
|
||||
cv2.imshow("Output", output_image)
|
||||
cv2.waitKey(0)
|
||||
@@ -0,0 +1,17 @@
|
||||
certifi==2025.4.26
|
||||
charset-normalizer==3.4.2
|
||||
coloredlogs==15.0.1
|
||||
flatbuffers==25.2.10
|
||||
humanfriendly==10.0
|
||||
idna==3.10
|
||||
mpmath==1.3.0
|
||||
numpy==2.0.2
|
||||
onnx==1.18.0
|
||||
onnxruntime==1.19.2
|
||||
packaging==25.0
|
||||
protobuf==6.31.1
|
||||
requests==2.32.4
|
||||
sympy==1.14.0
|
||||
typing-extensions==4.14.0
|
||||
urllib3==2.6.3
|
||||
opencv-python==4.11.0.86
|
||||
@@ -0,0 +1,137 @@
|
||||
# Ultralytics YOLO Interactive Object Tracking UI 🚀
|
||||
|
||||
A real-time [object detection](https://docs.ultralytics.com/tasks/detect/) and [tracking](https://docs.ultralytics.com/modes/track/) UI built with [Ultralytics YOLO11](https://github.com/ultralytics/ultralytics) and [OpenCV](https://opencv.org/), designed for interactive demos and seamless integration of tracking overlays. Whether you're just getting started with object tracking or looking to enhance it with additional features, this project provides a solid foundation.
|
||||
|
||||
https://github.com/user-attachments/assets/723e919e-555b-4cca-8e60-18e711d4f3b2
|
||||
|
||||
## ✨ Features
|
||||
|
||||
- Real-time object detection and visual tracking
|
||||
- Click-to-track any detected object
|
||||
- Scope lines and bold [bounding boxes](https://docs.ultralytics.com/usage/simple-utilities/#bounding-boxes) for active tracking
|
||||
- Dashed boxes for passive (non-tracked) objects
|
||||
- [Live terminal output](https://docs.ultralytics.com/guides/view-results-in-terminal/): object ID, label, [confidence](https://www.ultralytics.com/glossary/confidence), and center coordinates
|
||||
- Adjustable object tracking algorithms ([ByteTrack](https://docs.ultralytics.com/reference/trackers/byte_tracker/), [BoT-SORT](https://docs.ultralytics.com/reference/trackers/bot_sort/))
|
||||
- Supports:
|
||||
- [PyTorch](https://pytorch.org/) `.pt` models (for GPU devices like [NVIDIA Jetson](https://docs.ultralytics.com/guides/nvidia-jetson/) or [CUDA](https://developer.nvidia.com/cuda)-enabled desktops)
|
||||
- [NCNN](https://docs.ultralytics.com/integrations/ncnn/) `.param + .bin` models (for CPU-only devices like [Raspberry Pi](https://www.raspberrypi.org/) or ARM boards)
|
||||
|
||||
## 🏗️ Project Structure
|
||||
|
||||
```
|
||||
YOLO-Interactive-Tracking-UI/
|
||||
├── interactive_tracker.py # Main Python tracking UI script
|
||||
└── README.md # You're here!
|
||||
```
|
||||
|
||||
## 💻 Hardware & Model Compatibility
|
||||
|
||||
| Platform | Model Format | Example Model | GPU Acceleration | Notes |
|
||||
| ---------------- | ------------------ | -------------------- | ---------------- | ------------------------------- |
|
||||
| Raspberry Pi 4/5 | NCNN (.param/.bin) | `yolov8n_ncnn_model` | ❌ CPU only | Recommended format for Pi/ARM |
|
||||
| Jetson Nano | PyTorch (.pt) | `yolov8n.pt` | ✅ CUDA | Real-time performance possible |
|
||||
| Desktop w/ GPU | PyTorch (.pt) | `yolov8s.pt` | ✅ CUDA | Best performance |
|
||||
| CPU-only laptops | NCNN (.param/.bin) | `yolov8n_ncnn_model` | ❌ | Decent performance (~10–15 FPS) |
|
||||
|
||||
_Note: Performance may vary based on the specific hardware, model complexity, and input resolution._
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
### Basic Dependencies
|
||||
|
||||
Install the core `ultralytics` package:
|
||||
|
||||
```bash
|
||||
pip install ultralytics
|
||||
```
|
||||
|
||||
> **Tip:** Use a virtual environment like `venv` or [`conda`](https://docs.ultralytics.com/guides/conda-quickstart/) (recommended) to manage dependencies.
|
||||
|
||||
> **GPU Support:** Install PyTorch based on your system and CUDA version by following the official guide: [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/)
|
||||
|
||||
## 🚀 Quickstart
|
||||
|
||||
### Step 1: Download, Convert, or Specify Model
|
||||
|
||||
- For pre-trained Ultralytics YOLO [models](https://docs.ultralytics.com/models/) (e.g., `yolo11s.pt` or `yolov8s.pt`), simply specify the model name in the script parameters (`model_file`). These models will be automatically downloaded and cached. You can also manually download them from [Ultralytics Assets Releases](https://github.com/ultralytics/assets/releases) and place them in the project folder.
|
||||
- If you're using a custom-trained YOLO model, ensure the model file is in the project folder or provide its relative path.
|
||||
- For CPU-only devices, export your chosen model (e.g., `yolov8n.pt`) to the [NCNN format](https://docs.ultralytics.com/integrations/ncnn/) using the Ultralytics `export` mode.
|
||||
|
||||
- **Supported Formats:**
|
||||
- `yolo11s.pt` (for GPU with PyTorch)
|
||||
- `yolov8n_ncnn_model` (directory containing `.param` and `.bin` files for CPU with NCNN)
|
||||
|
||||
### Step 2: Configure the Script
|
||||
|
||||
Edit the global parameters at the top of `interactive_tracker.py`:
|
||||
|
||||
```python
|
||||
# --- Configuration ---
|
||||
enable_gpu = False # Set True if running with CUDA and PyTorch model
|
||||
model_file = "yolo11s.pt" # Path to model file (.pt for GPU, _ncnn_model dir for CPU)
|
||||
show_fps = True # Display current FPS in the top-left corner
|
||||
show_conf = False # Display confidence score for each detection
|
||||
save_video = False # Set True to save the output video stream
|
||||
video_output_path = "interactive_tracker_output.avi" # Output video file name
|
||||
|
||||
# --- Detection & Tracking Parameters ---
|
||||
conf = 0.3 # Minimum confidence threshold for object detection
|
||||
iou = 0.3 # IoU threshold for Non-Maximum Suppression (NMS)
|
||||
max_det = 20 # Maximum number of objects to detect per frame
|
||||
|
||||
tracker = "bytetrack.yaml" # Tracker configuration: 'bytetrack.yaml' or 'botsort.yaml'
|
||||
track_args = {
|
||||
"persist": True, # Keep track history across frames
|
||||
"verbose": False, # Suppress detailed tracker debug output
|
||||
}
|
||||
|
||||
window_name = "Ultralytics YOLO Interactive Tracking" # Name for the OpenCV display window
|
||||
# --- End Configuration ---
|
||||
```
|
||||
|
||||
- **`enable_gpu`**: Set to `True` if you have a CUDA-compatible GPU and are using a `.pt` model. Keep `False` for NCNN models or CPU-only execution.
|
||||
- **`model_file`**: Ensure this points to the correct model file or directory based on `enable_gpu`.
|
||||
- **`conf`**: Adjust the [confidence](https://www.ultralytics.com/glossary/confidence) threshold. Lower values detect more objects but may increase false positives.
|
||||
- **`iou`**: Set the [Intersection over Union (IoU)](https://www.ultralytics.com/glossary/intersection-over-union-iou) threshold for [Non-Maximum Suppression (NMS)](https://www.ultralytics.com/glossary/non-maximum-suppression-nms). Higher values allow more overlapping boxes.
|
||||
- **`tracker`**: Choose between available tracker configuration files ([ByteTrack](https://docs.ultralytics.com/reference/trackers/byte_tracker/), [BoT-SORT](https://docs.ultralytics.com/reference/trackers/bot_sort/)).
|
||||
|
||||
### Step 3: Run the Object Tracking
|
||||
|
||||
Execute the script from your terminal:
|
||||
|
||||
```bash
|
||||
python interactive_tracker.py
|
||||
```
|
||||
|
||||
### Controls
|
||||
|
||||
- 🖱️ **Left-click** on a detected object's bounding box to start tracking it.
|
||||
- 🔄 Press the **`c`** key to cancel the current tracking and select a new object.
|
||||
- ❌ Press the **`q`** key to quit the application.
|
||||
|
||||
### Saving Output Video (Optional)
|
||||
|
||||
If you want to record the tracking session, enable the `save_video` option in the configuration:
|
||||
|
||||
```python
|
||||
save_video = True # Enables video recording
|
||||
video_output_path = "output.avi" # Customize your output file name (e.g., .mp4, .avi)
|
||||
```
|
||||
|
||||
The video file will be saved in the project's working directory when you quit the application by pressing `q`.
|
||||
|
||||
## 👤 Author
|
||||
|
||||
- **Alireza**
|
||||
- [Connect on LinkedIn](https://www.linkedin.com/in/alireza787b)
|
||||
- Published: 2025-04-01
|
||||
|
||||
## 📜 License & Disclaimer
|
||||
|
||||
This project is released under the [AGPL-3.0 license](https://www.ultralytics.com/legal/agpl-3-0-software-license). For full licensing details, please refer to the [Ultralytics Licensing page](https://www.ultralytics.com/license).
|
||||
|
||||
This software is provided "as is" for educational and demonstration purposes. Use it responsibly and at your own risk. The author assumes no liability for misuse or unintended consequences.
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions, feedback, and bug reports are welcome! Feel free to open an issue or submit a pull request on the original repository if you have improvements or suggestions.
|
||||
@@ -0,0 +1,241 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import cv2
|
||||
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.utils import LOGGER
|
||||
from ultralytics.utils.plotting import Annotator, colors
|
||||
|
||||
enable_gpu = False # Set True if running with CUDA
|
||||
model_file = "yolo11s.pt" # Path to model file
|
||||
show_fps = True # If True, shows current FPS in top-left corner
|
||||
show_conf = False # Display or hide the confidence score
|
||||
save_video = False # Set True to save output video
|
||||
video_output_path = "interactive_tracker_output.avi" # Output video file name
|
||||
|
||||
|
||||
conf = 0.3 # Min confidence for object detection (lower = more detections, possibly more false positives)
|
||||
iou = 0.3 # IoU threshold for NMS (higher = less overlap allowed)
|
||||
max_det = 20 # Maximum objects per image (increase for crowded scenes)
|
||||
|
||||
tracker = "bytetrack.yaml" # Tracker config: 'bytetrack.yaml', 'botsort.yaml', etc.
|
||||
track_args = {
|
||||
"persist": True, # Keep frames history as a stream for continuous tracking
|
||||
"verbose": False, # Print debug info from tracker
|
||||
}
|
||||
|
||||
window_name = "Ultralytics YOLO Interactive Tracking" # Output window name
|
||||
|
||||
LOGGER.info("🚀 Initializing model...")
|
||||
if enable_gpu:
|
||||
LOGGER.info("Using GPU...")
|
||||
model = YOLO(model_file)
|
||||
model.to("cuda")
|
||||
else:
|
||||
LOGGER.info("Using CPU...")
|
||||
model = YOLO(model_file, task="detect")
|
||||
|
||||
classes = model.names # Store model class names
|
||||
|
||||
cap = cv2.VideoCapture(0) # Replace with video path if needed
|
||||
if not cap.isOpened():
|
||||
raise SystemError("Failed to open video source.")
|
||||
|
||||
vw = None # Initialized lazily after the first frame is read
|
||||
|
||||
selected_object_id = None
|
||||
selected_bbox = None
|
||||
selected_center = None
|
||||
latest_detections: list[list[float]] = []
|
||||
|
||||
|
||||
def get_center(x1: int, y1: int, x2: int, y2: int) -> tuple[int, int]:
|
||||
"""Calculate the center point of a bounding box.
|
||||
|
||||
Args:
|
||||
x1 (int): Top-left X coordinate.
|
||||
y1 (int): Top-left Y coordinate.
|
||||
x2 (int): Bottom-right X coordinate.
|
||||
y2 (int): Bottom-right Y coordinate.
|
||||
|
||||
Returns:
|
||||
center_x (int): X-coordinate of the center point.
|
||||
center_y (int): Y-coordinate of the center point.
|
||||
"""
|
||||
return (x1 + x2) // 2, (y1 + y2) // 2
|
||||
|
||||
|
||||
def extend_line_from_edge(mid_x: int, mid_y: int, direction: str, img_shape: tuple[int, int, int]) -> tuple[int, int]:
|
||||
"""Calculate the endpoint to extend a line from the center toward an image edge.
|
||||
|
||||
Args:
|
||||
mid_x (int): X-coordinate of the midpoint.
|
||||
mid_y (int): Y-coordinate of the midpoint.
|
||||
direction (str): Direction to extend ('left', 'right', 'up', 'down').
|
||||
img_shape (tuple[int, int, int]): Image shape in (height, width, channels).
|
||||
|
||||
Returns:
|
||||
end_x (int): X-coordinate of the endpoint.
|
||||
end_y (int): Y-coordinate of the endpoint.
|
||||
"""
|
||||
h, w = img_shape[:2]
|
||||
if direction == "down":
|
||||
return mid_x, h - 1
|
||||
elif direction == "left":
|
||||
return 0, mid_y
|
||||
elif direction == "right":
|
||||
return w - 1, mid_y
|
||||
elif direction == "up":
|
||||
return mid_x, 0
|
||||
else:
|
||||
return mid_x, mid_y
|
||||
|
||||
|
||||
def draw_tracking_scope(im, bbox: tuple, color: tuple) -> None:
|
||||
"""Draw tracking scope lines extending from the bounding box to image edges.
|
||||
|
||||
Args:
|
||||
im (np.ndarray): Image array to draw on.
|
||||
bbox (tuple): Bounding box coordinates (x1, y1, x2, y2).
|
||||
color (tuple): Color in BGR format for drawing.
|
||||
"""
|
||||
x1, y1, x2, y2 = bbox
|
||||
mid_top = ((x1 + x2) // 2, y1)
|
||||
mid_bottom = ((x1 + x2) // 2, y2)
|
||||
mid_left = (x1, (y1 + y2) // 2)
|
||||
mid_right = (x2, (y1 + y2) // 2)
|
||||
cv2.line(im, mid_top, extend_line_from_edge(*mid_top, "up", im.shape), color, 2)
|
||||
cv2.line(im, mid_bottom, extend_line_from_edge(*mid_bottom, "down", im.shape), color, 2)
|
||||
cv2.line(im, mid_left, extend_line_from_edge(*mid_left, "left", im.shape), color, 2)
|
||||
cv2.line(im, mid_right, extend_line_from_edge(*mid_right, "right", im.shape), color, 2)
|
||||
|
||||
|
||||
def click_event(event: int, x: int, y: int, flags: int, param) -> None:
|
||||
"""Handle mouse click events to select an object for focused tracking.
|
||||
|
||||
Args:
|
||||
event (int): OpenCV mouse event type.
|
||||
x (int): X-coordinate of the mouse event.
|
||||
y (int): Y-coordinate of the mouse event.
|
||||
flags (int): Any relevant flags passed by OpenCV.
|
||||
param (Any): Additional parameters (not used).
|
||||
"""
|
||||
global selected_object_id, latest_detections
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
if not latest_detections:
|
||||
return
|
||||
min_area = float("inf")
|
||||
best_match = None
|
||||
for track in latest_detections:
|
||||
if len(track) < 6:
|
||||
continue
|
||||
x1, y1, x2, y2 = map(int, track[:4])
|
||||
if x1 <= x <= x2 and y1 <= y <= y2:
|
||||
area = max(0, x2 - x1) * max(0, y2 - y1)
|
||||
if area < min_area:
|
||||
track_id = int(track[4]) if len(track) >= 7 else -1
|
||||
class_id = int(track[6]) if len(track) >= 7 else int(track[5])
|
||||
min_area = area
|
||||
best_match = (track_id, classes.get(class_id, str(class_id)))
|
||||
if best_match:
|
||||
selected_object_id, label = best_match
|
||||
LOGGER.info(f"Tracking started: {label} (ID {selected_object_id})")
|
||||
|
||||
|
||||
cv2.namedWindow(window_name)
|
||||
cv2.setMouseCallback(window_name, click_event)
|
||||
|
||||
fps_counter, fps_timer, fps_display = 0, time.time(), 0
|
||||
|
||||
while cap.isOpened():
|
||||
success, im = cap.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
results = model.track(im, conf=conf, iou=iou, max_det=max_det, tracker=tracker, **track_args)
|
||||
annotator = Annotator(im)
|
||||
detections = results[0].boxes.data if results[0].boxes is not None else []
|
||||
latest_detections = detections.cpu().tolist() if hasattr(detections, "cpu") else list(detections) # type: ignore[arg-type]
|
||||
detected_objects: list[str] = []
|
||||
for track in detections:
|
||||
track = track.tolist()
|
||||
if len(track) < 6:
|
||||
continue
|
||||
x1, y1, x2, y2 = map(int, track[:4])
|
||||
class_id = int(track[6]) if len(track) >= 7 else int(track[5])
|
||||
track_id = int(track[4]) if len(track) == 7 else -1
|
||||
color = colors(track_id, True)
|
||||
txt_color = annotator.get_txt_color(color)
|
||||
conf_score = float(track[5]) if len(track) >= 7 else 0.0
|
||||
class_name = classes.get(class_id, str(class_id))
|
||||
label = f"{class_name} ID {track_id}" + (f" ({conf_score:.2f})" if show_conf else "")
|
||||
center = get_center(x1, y1, x2, y2)
|
||||
detected_objects.append(f"{class_name}#{track_id}@{center[0]},{center[1]}")
|
||||
if track_id == selected_object_id:
|
||||
draw_tracking_scope(im, (x1, y1, x2, y2), color)
|
||||
cv2.circle(im, center, 6, color, -1)
|
||||
|
||||
# Pulsing circle for attention
|
||||
pulse_radius = 8 + int(4 * abs(time.time() % 1 - 0.5))
|
||||
cv2.circle(im, center, pulse_radius, color, 2)
|
||||
|
||||
annotator.box_label([x1, y1, x2, y2], label=f"ACTIVE: TRACK {track_id}", color=color)
|
||||
else:
|
||||
# Draw dashed box for other objects
|
||||
for i in range(x1, x2, 10):
|
||||
cv2.line(im, (i, y1), (i + 5, y1), color, 3)
|
||||
cv2.line(im, (i, y2), (i + 5, y2), color, 3)
|
||||
for i in range(y1, y2, 10):
|
||||
cv2.line(im, (x1, i), (x1, i + 5), color, 3)
|
||||
cv2.line(im, (x2, i), (x2, i + 5), color, 3)
|
||||
# Draw label text with background
|
||||
(tw, th), bl = cv2.getTextSize(label, 0, 0.7, 2)
|
||||
cv2.rectangle(im, (x1 + 5 - 5, y1 + 20 - th - 5), (x1 + 5 + tw + 5, y1 + 20 + bl), color, -1)
|
||||
cv2.putText(im, label, (x1 + 5, y1 + 20), 0, 0.7, txt_color, 1, cv2.LINE_AA)
|
||||
|
||||
if show_fps:
|
||||
fps_counter += 1
|
||||
if time.time() - fps_timer >= 1.0:
|
||||
fps_display = fps_counter
|
||||
fps_counter = 0
|
||||
fps_timer = time.time()
|
||||
|
||||
# Draw FPS text with background
|
||||
fps_text = f"FPS: {fps_display}"
|
||||
(tw, th), bl = cv2.getTextSize(fps_text, 0, 0.7, 2)
|
||||
cv2.rectangle(im, (10 - 5, 25 - th - 5), (10 + tw + 5, 25 + bl), (255, 255, 255), -1)
|
||||
cv2.putText(im, fps_text, (10, 25), 0, 0.7, (104, 31, 17), 1, cv2.LINE_AA)
|
||||
|
||||
if save_video and vw is None:
|
||||
h, w = im.shape[:2]
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) or 0
|
||||
fps = float(fps) if fps and fps > 0 else 30.0
|
||||
ext = video_output_path.lower()
|
||||
fourcc = cv2.VideoWriter_fourcc(*("MJPG" if ext.endswith(".avi") else "mp4v"))
|
||||
vw = cv2.VideoWriter(video_output_path, fourcc, fps, (w, h))
|
||||
|
||||
cv2.imshow(window_name, im)
|
||||
if save_video and vw is not None:
|
||||
vw.write(im)
|
||||
# Terminal logging
|
||||
LOGGER.info(
|
||||
f"Detected {len(detections)} object(s): {' | '.join(detected_objects)}"
|
||||
if detected_objects
|
||||
else f"Detected {len(detections)} object(s)."
|
||||
)
|
||||
|
||||
key = cv2.waitKey(1) & 0xFF
|
||||
if key == ord("q"):
|
||||
break
|
||||
elif key == ord("c"):
|
||||
LOGGER.info("Tracking reset.")
|
||||
selected_object_id = None
|
||||
|
||||
cap.release()
|
||||
if save_video and vw is not None:
|
||||
vw.release()
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,14 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
[package]
|
||||
name = "YOLO-ONNXRuntime-Rust"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
authors = ["Jamjamjon <xxyydzml@outlook.com>"]
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1.0.92"
|
||||
clap = "4.5.20"
|
||||
tracing = "0.1.40"
|
||||
tracing-subscriber = "0.3.18"
|
||||
usls = { version = "0.0.19", features = ["auto"] }
|
||||
@@ -0,0 +1,100 @@
|
||||
# YOLO-Series ONNXRuntime Rust Demo for Core YOLO Tasks
|
||||
|
||||
This repository provides a [Rust](https://rust-lang.org/) demo showcasing key [Ultralytics YOLO](https://docs.ultralytics.com/) series tasks such as [Classification](https://docs.ultralytics.com/tasks/classify/), [Segmentation](https://docs.ultralytics.com/tasks/segment/), [Detection](https://docs.ultralytics.com/tasks/detect/), [Pose Estimation](https://docs.ultralytics.com/tasks/pose/), and Oriented Bounding Box ([OBB](https://docs.ultralytics.com/tasks/obb/)) detection using the [ONNXRuntime](https://github.com/microsoft/onnxruntime). It supports various YOLO models (YOLOv5 through YOLO11) across multiple computer vision tasks.
|
||||
|
||||
## ✨ Introduction
|
||||
|
||||
- This example leverages the latest versions of both the [ONNX Runtime](https://onnxruntime.ai/) and popular YOLO models.
|
||||
- We utilize the [usls crate](https://github.com/jamjamjon/usls/tree/main) to streamline YOLO model inference in Rust, providing efficient data loading, visualization, and optimized inference performance. This allows developers to easily integrate state-of-the-art object detection into their Rust applications.
|
||||
|
||||
## 🚀 Features
|
||||
|
||||
- **Extensive Model Compatibility**: Supports a wide range of YOLO versions including [YOLOv5](https://docs.ultralytics.com/models/yolov5/), [YOLOv6](https://docs.ultralytics.com/models/yolov6/), [YOLOv7](https://docs.ultralytics.com/models/yolov7/), [YOLOv8](https://docs.ultralytics.com/models/yolov8/), [YOLOv9](https://docs.ultralytics.com/models/yolov9/), [YOLOv10](https://docs.ultralytics.com/models/yolov10/), [YOLO11](https://docs.ultralytics.com/models/yolo11/), [YOLO-World](https://docs.ultralytics.com/models/yolo-world/), [RT-DETR](https://docs.ultralytics.com/models/rtdetr/), and others.
|
||||
- **Versatile Task Coverage**: Includes examples for `Classification`, `Segmentation`, `Detection`, `Pose`, and `OBB`.
|
||||
- **Precision Flexibility**: Works seamlessly with `FP16` and `FP32` precision [ONNX models](https://docs.ultralytics.com/integrations/onnx/).
|
||||
- **Execution Providers**: Accelerated support for `CPU`, [CUDA](https://developer.nvidia.com/cuda-toolkit), [CoreML](https://developer.apple.com/documentation/coreml), and [TensorRT](https://docs.ultralytics.com/integrations/tensorrt/).
|
||||
- **Dynamic Input Shapes**: Dynamically adjusts to variable `batch`, `width`, and `height` dimensions for flexible model input.
|
||||
- **Flexible Data Loading**: The `DataLoader` component handles images, folders, videos, and real-time video streams.
|
||||
- **Real-Time Display and Video Export**: The `Viewer` provides real-time frame visualization and video export functions, similar to OpenCV’s `imshow()` and `imwrite()`.
|
||||
- **Enhanced Annotation and Visualization**: The `Annotator` facilitates comprehensive result rendering, supporting bounding boxes (HBB), oriented bounding boxes (OBB), polygons, masks, keypoints, and text labels.
|
||||
|
||||
## 🛠️ Setup Instructions
|
||||
|
||||
### 1. ONNXRuntime Linking
|
||||
|
||||
<details>
|
||||
<summary>You have two options to link the ONNXRuntime library:</summary>
|
||||
|
||||
- **Option 1: Manual Linking**
|
||||
- For detailed setup instructions, consult the [ONNX Runtime linking documentation](https://ort.pyke.io/setup/linking).
|
||||
- **Linux or macOS**:
|
||||
1. Download the appropriate ONNX Runtime package from the official [Releases page](https://github.com/microsoft/onnxruntime/releases).
|
||||
2. Set up the library path by exporting the `ORT_DYLIB_PATH` environment variable, pointing to the downloaded library file:
|
||||
```bash
|
||||
# Example path, replace with your actual path
|
||||
export ORT_DYLIB_PATH=/path/to/onnxruntime/lib/libonnxruntime.so.1.19.0
|
||||
```
|
||||
|
||||
- **Option 2: Automatic Download**
|
||||
- Use the `--features auto` flag with Cargo to let the build script handle downloading the library automatically:
|
||||
```bash
|
||||
cargo run -r --example yolo --features auto
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### 2. [Optional] Install CUDA, CuDNN, and TensorRT
|
||||
|
||||
- The CUDA execution provider requires [NVIDIA CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit) version `12.x`.
|
||||
- The TensorRT execution provider requires both CUDA `12.x` and [NVIDIA TensorRT](https://developer.nvidia.com/tensorrt) `10.x`. Ensure [cuDNN](https://developer.nvidia.com/cudnn) is also correctly installed.
|
||||
|
||||
### 3. [Optional] Install ffmpeg
|
||||
|
||||
To enable viewing video frames and saving video inferences, install the `rust-ffmpeg` crate's dependencies. Follow the instructions provided here:
|
||||
[https://github.com/zmwangx/rust-ffmpeg/wiki/Notes-on-building#dependencies](https://github.com/zmwangx/rust-ffmpeg/wiki/Notes-on-building#dependencies)
|
||||
|
||||
## ▶️ Get Started
|
||||
|
||||
Run the examples using Cargo. The `--` separates Cargo arguments from the example's arguments.
|
||||
|
||||
```bash
|
||||
# Run a custom model (e.g., YOLOv8 detection)
|
||||
cargo run -r -- --task detect --ver v8 --nc 6 --model path/to/your/model.onnx
|
||||
|
||||
# Classify examples
|
||||
cargo run -r -- --task classify --ver v5 --scale s --width 224 --height 224 --nc 1000 # YOLOv5 Classification
|
||||
cargo run -r -- --task classify --ver v8 --scale n --width 224 --height 224 --nc 1000 # YOLOv8 Classification
|
||||
cargo run -r -- --task classify --ver v11 --scale n --width 224 --height 224 --nc 1000 # YOLO11 Classification
|
||||
|
||||
# Detect examples
|
||||
cargo run -r -- --task detect --ver v5 --scale n # YOLOv5 Detection
|
||||
cargo run -r -- --task detect --ver v6 --scale n # YOLOv6 Detection
|
||||
cargo run -r -- --task detect --ver v7 --scale t # YOLOv7 Detection
|
||||
cargo run -r -- --task detect --ver v8 --scale n # YOLOv8 Detection
|
||||
cargo run -r -- --task detect --ver v9 --scale t # YOLOv9 Detection
|
||||
cargo run -r -- --task detect --ver v10 --scale n # YOLOv10 Detection
|
||||
cargo run -r -- --task detect --ver v11 --scale n # YOLO11 Detection
|
||||
cargo run -r -- --task detect --ver rtdetr --scale l # RT-DETR Detection
|
||||
|
||||
# Pose examples
|
||||
cargo run -r -- --task pose --ver v8 --scale n # YOLOv8-Pose Estimation
|
||||
cargo run -r -- --task pose --ver v11 --scale n # YOLO11-Pose Estimation
|
||||
|
||||
# Segment examples
|
||||
cargo run -r -- --task segment --ver v5 --scale n # YOLOv5-Segment
|
||||
cargo run -r -- --task segment --ver v8 --scale n # YOLOv8-Segment
|
||||
cargo run -r -- --task segment --ver v11 --scale n # YOLO11-Segment
|
||||
cargo run -r -- --task segment --ver v8 --model path/to/FastSAM-s-dyn-f16.onnx # FastSAM Segmentation
|
||||
|
||||
# OBB (Oriented Bounding Box) examples
|
||||
cargo run -r -- --ver v8 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLOv8-OBB
|
||||
cargo run -r -- --ver v11 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLO11-OBB
|
||||
```
|
||||
|
||||
**Use `cargo run -- --help` to see all available options.**
|
||||
|
||||
For more detailed information and advanced usage, please refer to the [usls-yolo example documentation](https://github.com/jamjamjon/usls/tree/main/examples/yolo).
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! If you'd like to improve this demo or add new features, please feel free to submit issues or pull requests on the repository. Your input helps make the Ultralytics ecosystem better for everyone. Check out the [Ultralytics Contribution Guide](https://docs.ultralytics.com/help/contributing/) for more details.
|
||||
@@ -0,0 +1,238 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
use anyhow::Result;
|
||||
use clap::Parser;
|
||||
|
||||
use usls::{
|
||||
models::YOLO, Annotator, DataLoader, Device, Options, Viewer, Vision, YOLOScale, YOLOTask,
|
||||
YOLOVersion, COCO_SKELETONS_16,
|
||||
};
|
||||
|
||||
#[derive(Parser, Clone)]
|
||||
#[command(author, version, about, long_about = None)]
|
||||
pub struct Args {
|
||||
/// Path to the ONNX model
|
||||
#[arg(long)]
|
||||
pub model: Option<String>,
|
||||
|
||||
/// Input source path
|
||||
#[arg(long, default_value_t = String::from("../../ultralytics/assets/bus.jpg"))]
|
||||
pub source: String,
|
||||
|
||||
/// YOLO Task
|
||||
#[arg(long, value_enum, default_value_t = YOLOTask::Detect)]
|
||||
pub task: YOLOTask,
|
||||
|
||||
/// YOLO Version
|
||||
#[arg(long, value_enum, default_value_t = YOLOVersion::V8)]
|
||||
pub ver: YOLOVersion,
|
||||
|
||||
/// YOLO Scale
|
||||
#[arg(long, value_enum, default_value_t = YOLOScale::N)]
|
||||
pub scale: YOLOScale,
|
||||
|
||||
/// Batch size
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub batch_size: usize,
|
||||
|
||||
/// Minimum input width
|
||||
#[arg(long, default_value_t = 224)]
|
||||
pub width_min: isize,
|
||||
|
||||
/// Input width
|
||||
#[arg(long, default_value_t = 640)]
|
||||
pub width: isize,
|
||||
|
||||
/// Maximum input width
|
||||
#[arg(long, default_value_t = 1024)]
|
||||
pub width_max: isize,
|
||||
|
||||
/// Minimum input height
|
||||
#[arg(long, default_value_t = 224)]
|
||||
pub height_min: isize,
|
||||
|
||||
/// Input height
|
||||
#[arg(long, default_value_t = 640)]
|
||||
pub height: isize,
|
||||
|
||||
/// Maximum input height
|
||||
#[arg(long, default_value_t = 1024)]
|
||||
pub height_max: isize,
|
||||
|
||||
/// Number of classes
|
||||
#[arg(long, default_value_t = 80)]
|
||||
pub nc: usize,
|
||||
|
||||
/// Class confidence
|
||||
#[arg(long)]
|
||||
pub confs: Vec<f32>,
|
||||
|
||||
/// Enable TensorRT support
|
||||
#[arg(long)]
|
||||
pub trt: bool,
|
||||
|
||||
/// Enable CUDA support
|
||||
#[arg(long)]
|
||||
pub cuda: bool,
|
||||
|
||||
/// Enable CoreML support
|
||||
#[arg(long)]
|
||||
pub coreml: bool,
|
||||
|
||||
/// Use TensorRT half precision
|
||||
#[arg(long)]
|
||||
pub half: bool,
|
||||
|
||||
/// Device ID to use
|
||||
#[arg(long, default_value_t = 0)]
|
||||
pub device_id: usize,
|
||||
|
||||
/// Enable performance profiling
|
||||
#[arg(long)]
|
||||
pub profile: bool,
|
||||
|
||||
/// Disable contour drawing, for saving time
|
||||
#[arg(long)]
|
||||
pub no_contours: bool,
|
||||
|
||||
/// Show result
|
||||
#[arg(long)]
|
||||
pub view: bool,
|
||||
|
||||
/// Do not save output
|
||||
#[arg(long)]
|
||||
pub nosave: bool,
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
|
||||
// logger
|
||||
if args.profile {
|
||||
tracing_subscriber::fmt()
|
||||
.with_max_level(tracing::Level::INFO)
|
||||
.init();
|
||||
}
|
||||
|
||||
// model path
|
||||
let path = match &args.model {
|
||||
None => format!(
|
||||
"yolo/{}-{}-{}.onnx",
|
||||
args.ver.name(),
|
||||
args.scale.name(),
|
||||
args.task.name()
|
||||
),
|
||||
Some(x) => x.to_string(),
|
||||
};
|
||||
|
||||
// saveout
|
||||
let saveout = match &args.model {
|
||||
None => format!(
|
||||
"{}-{}-{}",
|
||||
args.ver.name(),
|
||||
args.scale.name(),
|
||||
args.task.name()
|
||||
),
|
||||
Some(x) => {
|
||||
let p = std::path::PathBuf::from(&x);
|
||||
p.file_stem().unwrap().to_str().unwrap().to_string()
|
||||
}
|
||||
};
|
||||
|
||||
// device
|
||||
let device = if args.cuda {
|
||||
Device::Cuda(args.device_id)
|
||||
} else if args.trt {
|
||||
Device::Trt(args.device_id)
|
||||
} else if args.coreml {
|
||||
Device::CoreML(args.device_id)
|
||||
} else {
|
||||
Device::Cpu(args.device_id)
|
||||
};
|
||||
|
||||
// build options
|
||||
let options = Options::new()
|
||||
.with_model(&path)?
|
||||
.with_yolo_version(args.ver)
|
||||
.with_yolo_task(args.task)
|
||||
.with_device(device)
|
||||
.with_trt_fp16(args.half)
|
||||
.with_ixx(0, 0, (1, args.batch_size as _, 4).into())
|
||||
.with_ixx(0, 2, (args.height_min, args.height, args.height_max).into())
|
||||
.with_ixx(0, 3, (args.width_min, args.width, args.width_max).into())
|
||||
.with_confs(if args.confs.is_empty() {
|
||||
&[0.2, 0.15]
|
||||
} else {
|
||||
&args.confs
|
||||
})
|
||||
.with_nc(args.nc)
|
||||
.with_find_contours(!args.no_contours) // find contours or not
|
||||
// .with_names(&COCO_CLASS_NAMES_80) // detection class names
|
||||
// .with_names2(&COCO_KEYPOINTS_17) // keypoints class names
|
||||
// .exclude_classes(&[0])
|
||||
// .retain_classes(&[0, 5])
|
||||
.with_profile(args.profile);
|
||||
|
||||
// build model
|
||||
let mut model = YOLO::new(options)?;
|
||||
|
||||
// build dataloader
|
||||
let dl = DataLoader::new(&args.source)?
|
||||
.with_batch(model.batch() as _)
|
||||
.build()?;
|
||||
|
||||
// build annotator
|
||||
let annotator = Annotator::default()
|
||||
.with_skeletons(&COCO_SKELETONS_16)
|
||||
.without_masks(true) // no masks plotting when doing segment task
|
||||
.with_bboxes_thickness(3)
|
||||
.with_keypoints_name(false) // enable keypoints names
|
||||
.with_saveout_subs(&["YOLO"])
|
||||
.with_saveout(&saveout);
|
||||
|
||||
// build viewer
|
||||
let mut viewer = if args.view {
|
||||
Some(Viewer::new().with_delay(5).with_scale(1.).resizable(true))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// run & annotate
|
||||
for (xs, _paths) in dl {
|
||||
let ys = model.forward(&xs, args.profile)?;
|
||||
let images_plotted = annotator.plot(&xs, &ys, !args.nosave)?;
|
||||
|
||||
// show image
|
||||
match &mut viewer {
|
||||
Some(viewer) => viewer.imshow(&images_plotted)?,
|
||||
None => continue,
|
||||
}
|
||||
|
||||
// check out window and key event
|
||||
match &mut viewer {
|
||||
Some(viewer) => {
|
||||
if !viewer.is_open() || viewer.is_key_pressed(usls::Key::Escape) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
None => continue,
|
||||
}
|
||||
|
||||
// write video
|
||||
if !args.nosave {
|
||||
match &mut viewer {
|
||||
Some(viewer) => viewer.write_batch(&images_plotted)?,
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// finish video write
|
||||
if !args.nosave {
|
||||
if let Some(viewer) = &mut viewer {
|
||||
viewer.finish_write()?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
cmake_minimum_required(VERSION 3.5)
|
||||
|
||||
project(YOLO11TritonCPP VERSION 0.1)
|
||||
|
||||
set(CMAKE_INCLUDE_CURRENT_DIR ON)
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
|
||||
# Require external Triton client directory as a parameter
|
||||
if(NOT DEFINED TRITON_CLIENT_DIR)
|
||||
message(FATAL_ERROR "Please specify -DTRITON_CLIENT_DIR=/path/to/tritonclient")
|
||||
endif()
|
||||
|
||||
# Triton-related paths
|
||||
set(Protobuf_DIR "${TRITON_CLIENT_DIR}/protobuf/lib/cmake/protobuf")
|
||||
set(gRPC_DIR "${TRITON_CLIENT_DIR}/grpc")
|
||||
set(c-ares_DIR "${TRITON_CLIENT_DIR}/c-ares/lib/cmake/c-ares")
|
||||
set(TritonClient_DIR "${TRITON_CLIENT_DIR}/lib/cmake/TritonClient")
|
||||
set(TritonCommon_DIR "${TRITON_CLIENT_DIR}/lib/cmake/TritonCommon")
|
||||
|
||||
# Compiler optimizations
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mf16c -mavx2 -O3 -ffast-math -march=native")
|
||||
|
||||
# OpenCV setup
|
||||
find_package(OpenCV REQUIRED)
|
||||
include_directories(${OpenCV_INCLUDE_DIRS})
|
||||
|
||||
# Find Triton packages
|
||||
find_package(TritonCommon REQUIRED)
|
||||
find_package(TritonClient REQUIRED)
|
||||
|
||||
# Project source files
|
||||
set(PROJECT_SOURCES
|
||||
main.cpp
|
||||
inference.cpp
|
||||
inference.hpp
|
||||
)
|
||||
|
||||
# Define executable target
|
||||
add_executable(${PROJECT_NAME} ${PROJECT_SOURCES})
|
||||
|
||||
# Include directories
|
||||
target_include_directories(${PROJECT_NAME}
|
||||
PRIVATE
|
||||
${TRITON_CLIENT_DIR}/include
|
||||
${OpenCV_INCLUDE_DIRS}
|
||||
)
|
||||
|
||||
# Link directories
|
||||
target_link_directories(${PROJECT_NAME}
|
||||
PRIVATE
|
||||
${TRITON_CLIENT_DIR}/lib
|
||||
)
|
||||
|
||||
# Link libraries
|
||||
target_link_libraries(${PROJECT_NAME}
|
||||
PRIVATE
|
||||
${OpenCV_LIBS}
|
||||
grpcclient
|
||||
)
|
||||
110
algorithms/dms_yolo/code/examples/YOLO11-Triton-CPP/README.md
Normal file
110
algorithms/dms_yolo/code/examples/YOLO11-Triton-CPP/README.md
Normal file
@@ -0,0 +1,110 @@
|
||||
# YOLO11 Triton Inference Server C++ Client
|
||||
|
||||
[](https://github.com/ultralytics/ultralytics)
|
||||
[](https://github.com/triton-inference-server/server)
|
||||
|
||||
This example demonstrates how to perform object detection using Ultralytics YOLO11 models deployed on the NVIDIA Triton Inference Server. The implementation highlights efficient image preprocessing, FP16 (half-precision) data conversion, seamless communication with the Triton server via gRPC, and visualization of detection results with bounding boxes and confidence scores.
|
||||
|
||||
## ⚡ Features
|
||||
|
||||
- **High-Performance Inference**: Utilizes FP16 (half-precision) data format for optimized memory usage and accelerated inference.
|
||||
- **Non-Maximum Suppression (NMS)**: Removes duplicate detections to ensure precise object detection results.
|
||||
- **Seamless Triton Integration**: Communicates with the NVIDIA Triton Inference Server via gRPC for efficient and scalable model serving.
|
||||
- **Detection Visualization**: Annotates images with bounding boxes, class labels, and confidence scores for intuitive result interpretation.
|
||||
|
||||
## 🛠️ Dependencies
|
||||
|
||||
Ensure you have the following dependencies installed before proceeding:
|
||||
|
||||
| Dependency | Version | Description |
|
||||
| ----------------------- | ------- | --------------------------------------------- |
|
||||
| Triton Inference Server | 22.06 | Running with a deployed FP16 YOLO11 model |
|
||||
| Triton Client libraries | 2.23 | Required for communication with Triton Server |
|
||||
| C++ compiler | C++ 17+ | For compiling the C++ client application |
|
||||
| OpenCV library | 3.4.15 | For image processing and visualization |
|
||||
| CMake | 3.5+ | For building the project |
|
||||
|
||||
For more information on Triton, see the [NVIDIA Triton Inference Server documentation](https://github.com/triton-inference-server/server) and explore [model deployment options with Ultralytics](https://docs.ultralytics.com/guides/model-deployment-options/).
|
||||
|
||||
## 🏗️ Building the Project
|
||||
|
||||
1. **Install the Triton Client libraries:**
|
||||
|
||||
```bash
|
||||
wget https://github.com/triton-inference-server/server/releases/download/v2.23.0/v2.23.0_ubuntu2004.clients.tar.gz
|
||||
mkdir tritonclient
|
||||
tar -xvf v2.23.0_ubuntu2004.clients.tar.gz -C tritonclient
|
||||
rm -f v2.23.0_ubuntu2004.clients.tar.gz
|
||||
```
|
||||
|
||||
2. **Clone the Ultralytics repository:**
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ultralytics/ultralytics.git
|
||||
cd ultralytics/examples/YOLO11-Triton-CPP
|
||||
```
|
||||
|
||||
3. **Configure and build the project using CMake:**
|
||||
|
||||
```bash
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. -DTRITON_CLIENT_DIR=/path/to/tritonclient
|
||||
make
|
||||
```
|
||||
|
||||
For additional guidance on integrating Ultralytics YOLO models with various platforms, check out the [Ultralytics integrations documentation](https://docs.ultralytics.com/integrations/).
|
||||
|
||||
## 🚀 Usage
|
||||
|
||||
1. **Deploy your FP16 (half-precision) YOLO11 model on a Triton Inference Server.**
|
||||
Learn more about deploying models with [Ultralytics YOLO](https://docs.ultralytics.com/models/yolo11/).
|
||||
|
||||
2. **Run the YOLO11-Triton-CPP application:**
|
||||
|
||||
```bash
|
||||
./YOLO11TritonCPP
|
||||
```
|
||||
|
||||
By default, the application will:
|
||||
|
||||
- Connect to the Triton server at `localhost:8001`
|
||||
- Use the model named `yolo11` with version `1`
|
||||
- Process the image file `test.jpg`
|
||||
- Save detection results to `output.jpg`
|
||||
|
||||
For more on object detection workflows, see [Ultralytics object detection tasks](https://docs.ultralytics.com/tasks/detect/).
|
||||
|
||||
## ⚙️ Configuration
|
||||
|
||||
You can modify the following parameters in [main.cpp](main.cpp):
|
||||
|
||||
```cpp
|
||||
std::string triton_address = "localhost:8001";
|
||||
std::string model_name = "yolo11";
|
||||
std::string model_version = "1";
|
||||
std::string image_path = "test.jpg";
|
||||
std::string output_path = "output.jpg";
|
||||
std::vector<std::string> object_class_list = {"class1", "class2"};
|
||||
```
|
||||
|
||||
To learn more about configuring and customizing YOLO models, visit the [Ultralytics configuration guide](https://docs.ultralytics.com/usage/cfg/).
|
||||
|
||||
## 🌟 Contributors
|
||||
|
||||
Contributions are welcome! If you find any issues or have suggestions for improvements, please open an issue or submit a pull request on the [main Ultralytics repository](https://github.com/ultralytics/ultralytics).
|
||||
|
||||
- Ahmet Selim Demirel
|
||||
- Doğan Mehmet Başoğlu
|
||||
- Enes Uzun
|
||||
- Elif Cansu Ada
|
||||
- Mevlüt Ardıç
|
||||
- Serhat Karaca
|
||||
|
||||
[](https://github.com/ultralytics/ultralytics/graphs/contributors)
|
||||
|
||||
---
|
||||
|
||||
For more resources, explore the [Ultralytics documentation](https://docs.ultralytics.com/) and [Ultralytics blog](https://www.ultralytics.com/blog).
|
||||
|
||||
**We encourage your contributions to help improve this project.**
|
||||
@@ -0,0 +1,306 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include "inference.hpp"
|
||||
#include <immintrin.h>
|
||||
#include <iostream>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <cmath>
|
||||
|
||||
#define IOU_THRESHOLD 0.45
|
||||
|
||||
uint16_t float32_to_float16(float value) {
|
||||
__m128 input = _mm_set_ss(value);
|
||||
return _mm_cvtsi128_si32(_mm_cvtps_ph(input, 0));
|
||||
}
|
||||
|
||||
float float16_to_float32(uint16_t half) {
|
||||
uint16_t sign = (half & 0x8000) >> 15;
|
||||
uint16_t exponent = (half & 0x7C00) >> 10;
|
||||
uint16_t mantissa = (half & 0x03FF);
|
||||
|
||||
if (exponent == 0) {
|
||||
if (mantissa == 0) {
|
||||
return sign ? -0.0f : 0.0f;
|
||||
}
|
||||
return std::ldexp(mantissa / 1024.0f, -14) * (sign ? -1.0f : 1.0f);
|
||||
}
|
||||
else if (exponent == 31) {
|
||||
return mantissa ? NAN : (sign ? -INFINITY : INFINITY);
|
||||
}
|
||||
|
||||
float real_value = std::ldexp(1.0f + mantissa / 1024.0f, exponent - 15);
|
||||
return sign ? -real_value : real_value;
|
||||
}
|
||||
|
||||
|
||||
void Image::preprocess(cv::Mat* img, std::vector<uint16_t>& triton_data, int input_w, int input_h)
|
||||
{
|
||||
int w, h, x, y;
|
||||
float r_w = input_w / (img->cols*1.0);
|
||||
float r_h = input_h / (img->rows*1.0);
|
||||
if (r_h > r_w) {
|
||||
w = input_w;
|
||||
h = r_w * img->rows;
|
||||
x = 0;
|
||||
y = (input_h - h) / 2;
|
||||
} else {
|
||||
w = r_h * img->cols;
|
||||
h = input_h;
|
||||
x = (input_w - w) / 2;
|
||||
y = 0;
|
||||
}
|
||||
cv::Mat re(h, w, CV_8UC3);
|
||||
cv::resize(*img, re, re.size(), 0, 0, cv::INTER_LINEAR);
|
||||
cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(114, 114, 114));
|
||||
re.copyTo(out(cv::Rect(x, y, re.cols, re.rows)));
|
||||
cv::cvtColor(out, out, cv::COLOR_BGR2RGB);
|
||||
unsigned char* data = (unsigned char*)out.data;
|
||||
int step = out.step;
|
||||
|
||||
for (int yy = 0; yy < input_h; ++yy)
|
||||
{
|
||||
for (int kk = 0; kk < 3; ++kk)
|
||||
{
|
||||
for (int xx = 0; xx < input_w; ++xx)
|
||||
{
|
||||
float temp_f = data[yy * step + xx * 3 + kk] / 255.0f;
|
||||
|
||||
triton_data[kk * input_w * input_h + yy * input_w + xx] = float32_to_float16(temp_f);
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int getDetectionsFromTritonRawData(std::vector<float>& detection_results, std::vector<struct detection_struct> &detections, std::vector<std::string>& object_class_list, float confidence_threshold, int image_width, int image_height)
|
||||
{
|
||||
const size_t shape[3] = {1, object_class_list.size()+4, 8400};
|
||||
std::vector<BoundingBox> boxes;
|
||||
for (size_t i = 0; i < shape[2]; i++)
|
||||
{
|
||||
|
||||
int x = int(detection_results[0 * shape[2] + i]);
|
||||
int y = int(detection_results[1 * shape[2] + i]);
|
||||
int w = int(detection_results[2 * shape[2] + i]);
|
||||
int h = int(detection_results[3 * shape[2] + i]);
|
||||
|
||||
for(size_t j =0 ; j < object_class_list.size(); j++)
|
||||
{
|
||||
if (detection_results[(4+j) * shape[2] + i] > 0.01)
|
||||
{
|
||||
BoundingBox box;
|
||||
box.x = static_cast<float>(x);
|
||||
box.y = static_cast<float>(y);
|
||||
box.w = static_cast<float>(w);
|
||||
box.h = static_cast<float>(h);
|
||||
box.score = detection_results[(4 + j) * shape[2] + i];
|
||||
box.class_id = j;
|
||||
boxes.push_back(box);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
auto nms_boxes = NMS(boxes, IOU_THRESHOLD);
|
||||
detections.clear();
|
||||
if(nms_boxes.size()==0)
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
float scale_x = 0.0;
|
||||
float scale_y = 0.0;
|
||||
|
||||
int x1,y1 = 0;
|
||||
int x2,y2 = 0;
|
||||
|
||||
float shift_factor_x = 0.6;
|
||||
float shift_factor_y = 0.5;
|
||||
|
||||
int offset_shift = (image_width/640.0f)*10;
|
||||
|
||||
if (image_width<=640)
|
||||
{
|
||||
scale_x = static_cast<float>(image_width - 640.0f ) * 0.5 ;
|
||||
scale_y = static_cast<float>(image_height - 640.0f) * 0.5 ;
|
||||
}
|
||||
for (size_t i = 0; i < nms_boxes.size(); ++i)
|
||||
{
|
||||
if (nms_boxes[i].score< confidence_threshold)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
struct detection_struct tespit_yapi ;
|
||||
detections.push_back(tespit_yapi);
|
||||
detections[detections.size() - 1].confidence_score = nms_boxes[i].score;
|
||||
if (image_width==640)
|
||||
{
|
||||
scale_x = static_cast<float>(image_width - 640.0f ) * 0.5 ;
|
||||
scale_y = static_cast<float>(image_height - 640.0f) * 0.5 ;
|
||||
x1 = static_cast<int>((nms_boxes[i].x - nms_boxes[i].w/2) + scale_x);
|
||||
y1 = static_cast<int>((nms_boxes[i].y - nms_boxes[i].h/2) + scale_y) ;
|
||||
x2 = static_cast<int>((nms_boxes[i].x + nms_boxes[i].w/2) + scale_x);
|
||||
y2 = static_cast<int>((nms_boxes[i].y + nms_boxes[i].h/2) + scale_y);
|
||||
}
|
||||
else if(image_width>=1080)
|
||||
{
|
||||
x1 = static_cast<int>((nms_boxes[i].x - nms_boxes[i].w/2) * (image_width/640) );
|
||||
y1 = static_cast<int>((nms_boxes[i].y - nms_boxes[i].h/2) * (image_width/640) - ((image_width - image_height) / 2.0)) ;
|
||||
x2 = static_cast<int>((nms_boxes[i].x + nms_boxes[i].w/2) * (image_width/640) );
|
||||
y2 = static_cast<int>((nms_boxes[i].y + nms_boxes[i].h/2) * (image_width/640)- ((image_width - image_height) / 2.0));
|
||||
}
|
||||
|
||||
float x_center, y_center, width, height;
|
||||
x_center = (x1 + x2) / 2.0f;
|
||||
y_center = (y1 + y2) / 2.0f;
|
||||
width = x2 - x1;
|
||||
height = y2 - y1;
|
||||
detections[detections.size() - 1].bbox.x = x_center - width/2 ;
|
||||
detections[detections.size() - 1].bbox.y = y_center - height/2;
|
||||
detections[detections.size() - 1].bbox.width = width ;
|
||||
detections[detections.size() - 1].bbox.height = height;
|
||||
if (detections[detections.size() - 1].bbox.x <= 0)
|
||||
detections[detections.size() - 1].bbox.x = offset_shift;
|
||||
if (detections[detections.size() - 1].bbox.y <= 0)
|
||||
detections[detections.size() - 1].bbox.y = offset_shift;
|
||||
if (detections[detections.size() - 1].bbox.x + detections[detections.size() - 1].bbox.width >= image_width)
|
||||
{
|
||||
detections[detections.size() - 1].bbox.width -= detections[detections.size() - 1].bbox.x + detections[detections.size() - 1].bbox.width - image_width + offset_shift ;
|
||||
}
|
||||
if (detections[detections.size() - 1].bbox.y + detections[detections.size() - 1].bbox.height >= image_height)
|
||||
{
|
||||
detections[detections.size() - 1].bbox.height -= detections[detections.size() - 1].bbox.y + detections[detections.size() - 1].bbox.height - image_height + offset_shift ;
|
||||
}
|
||||
detections[detections.size() - 1].name = object_class_list[nms_boxes[i].class_id];
|
||||
detections[detections.size() - 1].class_id = nms_boxes[i].class_id;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
float IoU(const BoundingBox& box1, const BoundingBox& box2) {
|
||||
float x1_min = box1.x - box1.w / 2.0f;
|
||||
float y1_min = box1.y - box1.h / 2.0f;
|
||||
float x1_max = box1.x + box1.w / 2.0f;
|
||||
float y1_max = box1.y + box1.h / 2.0f;
|
||||
|
||||
float x2_min = box2.x - box2.w / 2.0f;
|
||||
float y2_min = box2.y - box2.h / 2.0f;
|
||||
float x2_max = box2.x + box2.w / 2.0f;
|
||||
float y2_max = box2.y + box2.h / 2.0f;
|
||||
|
||||
float inter_x_min = std::max(x1_min, x2_min);
|
||||
float inter_y_min = std::max(y1_min, y2_min);
|
||||
float inter_x_max = std::min(x1_max, x2_max);
|
||||
float inter_y_max = std::min(y1_max, y2_max);
|
||||
|
||||
float inter_width = inter_x_max - inter_x_min;
|
||||
float inter_height = inter_y_max - inter_y_min;
|
||||
|
||||
if (inter_width <= 0 || inter_height <= 0)
|
||||
return 0.0f;
|
||||
|
||||
float inter_area = inter_width * inter_height;
|
||||
float area1 = (x1_max - x1_min) * (y1_max - y1_min);
|
||||
float area2 = (x2_max - x2_min) * (y2_max - y2_min);
|
||||
float union_area = area1 + area2 - inter_area;
|
||||
|
||||
return inter_area / union_area;
|
||||
}
|
||||
|
||||
std::vector<BoundingBox> NMS(const std::vector<BoundingBox>& boxes, float iou_threshold) {
|
||||
|
||||
std::vector<BoundingBox> result;
|
||||
|
||||
std::vector<BoundingBox> sorted_boxes = boxes;
|
||||
std::sort(sorted_boxes.begin(), sorted_boxes.end(), [](const BoundingBox& a, const BoundingBox& b) {
|
||||
return a.score > b.score;
|
||||
});
|
||||
|
||||
std::vector<bool> suppressed(sorted_boxes.size(), false);
|
||||
|
||||
for (size_t i = 0; i < sorted_boxes.size(); ++i) {
|
||||
if (suppressed[i])
|
||||
continue;
|
||||
|
||||
result.push_back(sorted_boxes[i]);
|
||||
|
||||
for (size_t j = i + 1; j < sorted_boxes.size(); ++j) {
|
||||
if (suppressed[j])
|
||||
continue;
|
||||
if (IoU(sorted_boxes[i], sorted_boxes[j]) > iou_threshold) {
|
||||
suppressed[j] = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
TritonCommunication::TritonCommunication(std::string triton_address, std::string model_name, std::string model_version, int image_channel, int image_width, int image_height, int class_count) : options(model_name)
|
||||
{
|
||||
triton::client::Error err;
|
||||
this->triton_url = triton_address;
|
||||
|
||||
this->options.model_version_ = model_version;
|
||||
|
||||
this->shape = {1, image_channel, image_width, image_height};
|
||||
this->input_byte_size = image_channel * image_width * image_height * sizeof(uint16_t) ;
|
||||
this->output_byte_size = (class_count + 4) * 8400 * sizeof(uint16_t);
|
||||
|
||||
err = tc::InferenceServerGrpcClient::Create(&(this->client), this->triton_url);
|
||||
if (!err.IsOk()) {
|
||||
std::cout<< "Create grpc client error:"<<err.Message()<<std::endl;
|
||||
}
|
||||
bool live;
|
||||
err = client->IsServerReady(&live);
|
||||
if (!err.IsOk() || !live) {
|
||||
std::cout<< "Triton server is not live !"<<std::endl;
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
bool model_ready;
|
||||
err = client->IsModelReady(&model_ready,model_name,model_version);
|
||||
if (!err.IsOk() || !model_ready) {
|
||||
std::cerr << "Model:[" << model_name << "] has not been deployed on Triton Server. Triton Server Address:["<<triton_address <<"]"<<std::endl;
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
std::cout<<"Triton server is LIVE and model is READY!"<<std::endl;
|
||||
}
|
||||
|
||||
void TritonCommunication::infer(uint16_t* image_data)
|
||||
{
|
||||
size_t num_elements = output_byte_size / sizeof(uint16_t);
|
||||
tc::Error err;
|
||||
tc::InferInput* input0;
|
||||
|
||||
err = tc::InferInput::Create(&input0, "images", shape, "FP16"); // FP16 is the data type of the input tensor.
|
||||
std::shared_ptr<tc::InferInput> input0_ptr;
|
||||
input0_ptr.reset(input0);
|
||||
|
||||
err = input0_ptr->AppendRaw((const uint8_t*)image_data, this->input_byte_size);
|
||||
|
||||
std::vector<tc::InferInput*> inputs = {input0_ptr.get()};
|
||||
|
||||
tc::InferResult* results;
|
||||
err = client->Infer(&results, options, inputs);
|
||||
results_ptr.reset(results);
|
||||
|
||||
float *output0_data;
|
||||
size_t output0_byte_size;
|
||||
std::vector<uint16_t> result_fp16_raw_data;
|
||||
|
||||
results->RawData("output0", (const uint8_t**)&output0_data, &output0_byte_size); // output0 is a specific name for the output tensor.
|
||||
result_fp16_raw_data.resize(output0_byte_size/sizeof(uint16_t));
|
||||
|
||||
std::memcpy(result_fp16_raw_data.data(), output0_data, output0_byte_size);
|
||||
std::vector<float> float32_data(num_elements);
|
||||
|
||||
for (size_t i = 0; i < num_elements; i++) {
|
||||
float32_data[i] = float16_to_float32(result_fp16_raw_data[i]);
|
||||
}
|
||||
|
||||
output_raw_data = float32_data;
|
||||
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include "grpc_client.h"
|
||||
|
||||
class Image
|
||||
{
|
||||
public:
|
||||
Image() = default;
|
||||
static void preprocess(cv::Mat* image, std::vector<uint16_t>& triton_data, int input_w, int input_h);
|
||||
};
|
||||
|
||||
struct struct_yolo_output
|
||||
{
|
||||
std::vector<int> num_dets, det_classes;
|
||||
std::vector<float> det_boxes, det_scores;
|
||||
};
|
||||
|
||||
struct BoundingBox {
|
||||
float x, y, w, h;
|
||||
float score;
|
||||
int class_id;
|
||||
};
|
||||
|
||||
struct detection_struct
|
||||
{
|
||||
cv::Rect bbox;
|
||||
int class_id;
|
||||
std::string name;
|
||||
double confidence_score;
|
||||
};
|
||||
|
||||
// C-compatible declarations
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
int getDetectionsFromTritonRawData(
|
||||
std::vector<float>& detection_results,
|
||||
std::vector<struct detection_struct>& tespitler,
|
||||
std::vector<std::string>& object_class_list,
|
||||
float confidence_threshold,
|
||||
int image_width,
|
||||
int image_height
|
||||
);
|
||||
|
||||
std::vector<BoundingBox> NMS(const std::vector<BoundingBox>& boxes, float iou_threshold);
|
||||
float IoU(const BoundingBox& box1, const BoundingBox& box2);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
namespace tc = triton::client;
|
||||
|
||||
class TritonCommunication
|
||||
{
|
||||
private:
|
||||
std::unique_ptr<tc::InferenceServerGrpcClient> client;
|
||||
std::string triton_url;
|
||||
std::vector<int64_t> shape;
|
||||
tc::InferOptions options;
|
||||
size_t input_byte_size;
|
||||
size_t output_byte_size;
|
||||
std::shared_ptr<tc::InferResult> results_ptr;
|
||||
|
||||
public:
|
||||
std::vector<float> output_raw_data;
|
||||
|
||||
TritonCommunication(std::string triton_address,
|
||||
std::string model_name,
|
||||
std::string model_version,
|
||||
int image_channel,
|
||||
int image_width,
|
||||
int image_height,
|
||||
int class_count);
|
||||
|
||||
void infer(uint16_t* triton_data);
|
||||
};
|
||||
76
algorithms/dms_yolo/code/examples/YOLO11-Triton-CPP/main.cpp
Normal file
76
algorithms/dms_yolo/code/examples/YOLO11-Triton-CPP/main.cpp
Normal file
@@ -0,0 +1,76 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include "inference.hpp"
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <sstream>
|
||||
#include <iomanip>
|
||||
#include <cstdint>
|
||||
#include <chrono>
|
||||
|
||||
#define MODEL_INPUT_IMAGE_WIDTH 640
|
||||
#define MODEL_INPUT_IMAGE_HEIGHT 640
|
||||
#define NETWORK_THRESHOLD 0.50
|
||||
#define IMAGE_CHANNEL 3
|
||||
|
||||
double get_time_since_epoch_millis()
|
||||
{
|
||||
using namespace std::chrono;
|
||||
auto now = system_clock::now();
|
||||
auto duration = now.time_since_epoch();
|
||||
return duration_cast<microseconds>(duration).count() / 1000.0;
|
||||
}
|
||||
|
||||
int main(int argc, char *argv[])
|
||||
{
|
||||
std::string triton_address= "localhost:8001";
|
||||
std::string model_name= "yolo11";
|
||||
std::string model_version= "1";
|
||||
std::string image_path = "test.jpg";
|
||||
std::string output_path = "output.jpg";
|
||||
std::vector<std::string> object_class_list = {"class1", "class2"};
|
||||
|
||||
std::vector<uint16_t> triton_request_data;
|
||||
triton_request_data.resize(IMAGE_CHANNEL*MODEL_INPUT_IMAGE_WIDTH*MODEL_INPUT_IMAGE_HEIGHT);
|
||||
std::vector<struct detection_struct> detections;
|
||||
|
||||
std::shared_ptr<TritonCommunication> triton_communication = std::make_shared<TritonCommunication>(triton_address, model_name, model_version, IMAGE_CHANNEL, MODEL_INPUT_IMAGE_WIDTH, MODEL_INPUT_IMAGE_HEIGHT,object_class_list.size());
|
||||
|
||||
cv::Mat frame = cv::imread(image_path);
|
||||
if (frame.empty())
|
||||
{
|
||||
std::cerr << "Image couldn't read: " << image_path << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
int image_width = frame.cols;
|
||||
int image_height = frame.rows;
|
||||
|
||||
double preprocess_time = get_time_since_epoch_millis();
|
||||
Image::preprocess(&frame, triton_request_data, MODEL_INPUT_IMAGE_WIDTH, MODEL_INPUT_IMAGE_HEIGHT);
|
||||
std::cout << "Preprocess time : " << (get_time_since_epoch_millis() - preprocess_time)<< " millisecond."<< std::endl;
|
||||
|
||||
double infer_time = get_time_since_epoch_millis();
|
||||
triton_communication->infer(triton_request_data.data());
|
||||
std::cout << "Triton Server execute time : " << (get_time_since_epoch_millis() - infer_time) << " millisecond." << std::endl;
|
||||
|
||||
getDetectionsFromTritonRawData(triton_communication->output_raw_data, detections, object_class_list, NETWORK_THRESHOLD, image_width, image_height);
|
||||
|
||||
for (int i = 0; i < detections.size(); i++)
|
||||
{
|
||||
std::ostringstream oss;
|
||||
oss << detections[i].name << " "
|
||||
<< std::fixed << std::setprecision(2)
|
||||
<< detections[i].confidence_score;
|
||||
|
||||
cv::rectangle(frame, detections[i].bbox, cv::Scalar(255, 0, 0), 2);
|
||||
cv::putText(frame, oss.str(), cv::Point((detections[i].bbox.x), (detections[i].bbox.y - 5)), cv::FONT_HERSHEY_DUPLEX, ((frame.cols / 640.0f) * 0.35), cv::Scalar(0, 0, 0), (int)(frame.cols / 640.0f) + 1);
|
||||
cv::putText(frame, oss.str(), cv::Point((detections[i].bbox.x), (detections[i].bbox.y - 5)), cv::FONT_HERSHEY_DUPLEX, ((frame.cols / 640.0f) * 0.35), cv::Scalar(0xFF, 0xFF, 0xFF), (int)(frame.cols / 640.0f));
|
||||
}
|
||||
|
||||
cv::imwrite(output_path, frame);
|
||||
std::cout << "Result image saved!"<< std::endl;
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
# Zero-Shot Action Recognition with Ultralytics YOLO11 (Inference on Video)
|
||||
|
||||
Action recognition is a [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) technique used to identify and classify actions performed by individuals in a video. This process enables more advanced analyses when multiple actions are considered. Using models like [Ultralytics YOLO11](https://docs.ultralytics.com/models/yolo11/), actions can be detected and classified in real time. This system leverages **zero-shot learning**, allowing it to recognize actions it wasn't explicitly trained on by using descriptive labels. Learn more about zero-shot concepts on [Wikipedia](https://en.wikipedia.org/wiki/Zero-shot_learning).
|
||||
|
||||
The system can be customized to recognize specific actions based on the user's preferences and requirements by providing different text labels.
|
||||
|
||||
## 🎬 Table of Contents
|
||||
|
||||
- [Step 1: Install the Required Libraries](#step-1-install-the-required-libraries)
|
||||
- [Step 2: Run Action Recognition Using Ultralytics YOLO11](#step-2-run-action-recognition-using-ultralytics-yolo11)
|
||||
- [Usage Options](#usage-options)
|
||||
- [FAQ](#faq)
|
||||
|
||||
## ⚙️ Step 1: Install the Required Libraries
|
||||
|
||||
Clone the [Ultralytics GitHub repository](https://github.com/ultralytics/ultralytics) using [Git](https://git-scm.com/), install dependencies with [pip](https://pip.pypa.io/en/stable/), and navigate (`cd`) to this local directory for the commands in Step 2.
|
||||
|
||||
```bash
|
||||
# Clone ultralytics repo
|
||||
git clone https://github.com/ultralytics/ultralytics
|
||||
|
||||
# cd to local directory
|
||||
cd ultralytics/examples/YOLOv8-Action-Recognition
|
||||
|
||||
# Install dependencies using Python's package manager
|
||||
pip install -U -r requirements.txt
|
||||
```
|
||||
|
||||
## 🚀 Step 2: Run Action Recognition Using Ultralytics YOLO11
|
||||
|
||||
Here are the basic commands for running inference:
|
||||
|
||||
### Note
|
||||
|
||||
The action recognition model will automatically perform [object detection](https://www.ultralytics.com/glossary/object-detection) and [tracking](https://docs.ultralytics.com/modes/track/) for people in the video, and classify their actions based on the specified labels. The results will be displayed in real-time on the video output. You can customize the action labels by modifying the `--labels` argument when running the [Python](https://www.python.org/) script. This utilizes a video classifier model, often sourced from platforms like [Hugging Face Models](https://huggingface.co/models).
|
||||
|
||||
```bash
|
||||
# Quick start with default video and labels
|
||||
python action_recognition.py
|
||||
|
||||
# Basic usage with a YouTube video and custom labels
|
||||
python action_recognition.py --source "https://www.youtube.com/watch?v=dQw4w9WgXcQ" --labels "dancing" "singing a song"
|
||||
|
||||
# Use a local video file
|
||||
python action_recognition.py --source path/to/video.mp4
|
||||
|
||||
# Use a medium-sized YOLO11 model for potentially better detector performance
|
||||
python action_recognition.py --weights yolo11m.pt
|
||||
|
||||
# Run inference on the CPU instead of GPU
|
||||
python action_recognition.py --device cpu
|
||||
|
||||
# Use a different video classifier model from TorchVision
|
||||
python action_recognition.py --video-classifier-model "s3d"
|
||||
|
||||
# Use FP16 (half-precision) for faster inference (only for HuggingFace models)
|
||||
python action_recognition.py --fp16
|
||||
|
||||
# Export the output video with recognized actions to an mp4 file
|
||||
python action_recognition.py --output-path output.mp4
|
||||
|
||||
# Combine multiple options: specific YouTube source, GPU device 0, specific HuggingFace model, custom labels, and FP16
|
||||
python action_recognition.py --source "https://www.youtube.com/watch?v=dQw4w9WgXcQ" --device 0 --video-classifier-model "microsoft/xclip-base-patch32" --labels "dancing" "singing a song" --fp16
|
||||
```
|
||||
|
||||
## 🛠️ Usage Options
|
||||
|
||||
- `--weights`: Path to the YOLO [model weights](https://www.ultralytics.com/glossary/model-weights) file (default: `"yolo11n.pt"`). You can choose other models like `yolo11s.pt`, `yolo11m.pt`, etc.
|
||||
- `--device`: CUDA device identifier (e.g., `0` or `0,1,2,3`) or `cpu` to run on the [CPU](https://www.ultralytics.com/glossary/cpu) (default: auto-detects available [GPU](https://www.ultralytics.com/glossary/gpu-graphics-processing-unit)).
|
||||
- `--source`: Path to a local video file or a YouTube URL (default: "[rickroll](https://www.youtube.com/watch?v=dQw4w9WgXcQ)").
|
||||
- `--output-path`: Path to save the output video file (default: `output_video.mp4`).
|
||||
- `--crop-margin-percentage`: Percentage of margin to add around detected objects before cropping for classification (default: `10`).
|
||||
- `--num-video-sequence-samples`: Number of video frames sampled from a sequence to feed into the classifier (default: `8`).
|
||||
- `--skip-frame`: Number of frames to skip between detections to speed up processing (default: `2`).
|
||||
- `--video-cls-overlap-ratio`: Overlap ratio between consecutive video sequences sent for classification (default: `0.25`).
|
||||
- `--fp16`: Use [FP16 (half-precision)](https://www.ultralytics.com/glossary/half-precision) for inference, potentially speeding it up on compatible hardware (only applicable to Hugging Face models).
|
||||
- `--video-classifier-model`: Name or path of the video classifier model (default: `"microsoft/xclip-base-patch32"`). Can be a Hugging Face model name or a [TorchVision model](https://docs.pytorch.org/vision/stable/models.html) name.
|
||||
- `--labels`: A list of text labels for zero-shot video classification (default: `["dancing", "singing a song"]`).
|
||||
|
||||
## 🤔 FAQ
|
||||
|
||||
### 1. What Does Action Recognition Involve?
|
||||
|
||||
Action recognition is a computational method used to identify and classify actions or activities performed by individuals in recorded video or real-time streams. This technique is widely used in video analysis, surveillance, and human-computer interaction, enabling the detection and understanding of human behaviors based on their motion patterns and context. It often combines [object tracking](https://www.ultralytics.com/glossary/object-tracking) with classification. Explore more on [video classification research](https://arxiv.org/).
|
||||
|
||||
### 2. Are Custom Action Labels Supported by Action Recognition?
|
||||
|
||||
Yes, custom action labels are supported. The `action_recognition.py` script allows users to specify their own custom labels for **zero-shot video classification**. This is done using the `--labels` argument. For example:
|
||||
|
||||
```bash
|
||||
python action_recognition.py --source https://www.youtube.com/watch?v=dQw4w9WgXcQ --labels "walking" "running" "jumping"
|
||||
```
|
||||
|
||||
You can adjust these labels to match the specific actions you want the system to recognize in your video. The system will then attempt to classify detected actions based on these custom labels using its understanding derived from large datasets.
|
||||
|
||||
Additionally, you can choose between different video classification models:
|
||||
|
||||
1. **Hugging Face Models**: You can use any compatible video classification model available on Hugging Face Hub. The default is:
|
||||
- `"microsoft/xclip-base-patch32"`
|
||||
2. **TorchVision Models**: These models do not support zero-shot classification with custom text labels but offer pre-trained classification capabilities. Options include:
|
||||
- `"s3d"`
|
||||
- `"r3d_18"`
|
||||
- `"swin3d_t"`
|
||||
- `"swin3d_b"`
|
||||
- `"mvit_v1_b"`
|
||||
- `"mvit_v2_s"`
|
||||
|
||||
### 3. Why Combine Action Recognition with YOLO11?
|
||||
|
||||
[Ultralytics YOLO11](https://docs.ultralytics.com/models/yolo11/) excels at fast and accurate [object detection](https://docs.ultralytics.com/tasks/detect/) and tracking in video streams. Combining it with action recognition allows the system not only to locate individuals (using YOLO11's detection capabilities) but also to understand _what_ they are doing. This synergy provides a richer analysis of video content, crucial for applications like automated surveillance, sports analytics, or human-robot interaction. See our blog post on [object detection and tracking](https://www.ultralytics.com/blog/object-detection-and-tracking-with-ultralytics-yolov8).
|
||||
|
||||
### 4. Can I Employ Other YOLO Versions?
|
||||
|
||||
Certainly! While this example defaults to `yolo11n.pt`, you have the flexibility to specify different Ultralytics YOLO model weights using the `--weights` option. For instance, you could use `yolo11s.pt`, `yolo11m.pt`, `yolo11l.pt`, or `yolo11x.pt` for potentially higher detection accuracy at the cost of inference speed. You can even use models trained for other tasks if applicable, though detection models are standard here. Check the [Ultralytics documentation](https://docs.ultralytics.com/) for available models and their performance metrics.
|
||||
|
||||
---
|
||||
|
||||
We hope this guide helps you implement zero-shot action recognition using Ultralytics YOLO11! Feel free to explore the code and experiment with different options. If you encounter issues or have suggestions, please consider contributing by opening an issue or pull request on the [GitHub repository](https://github.com/ultralytics/ultralytics). See our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for more details.
|
||||
@@ -0,0 +1,511 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import AutoModel, AutoProcessor
|
||||
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.data.loaders import get_best_youtube_url
|
||||
from ultralytics.utils.plotting import Annotator
|
||||
from ultralytics.utils.torch_utils import select_device
|
||||
|
||||
|
||||
class TorchVisionVideoClassifier:
|
||||
"""Video classifier using pretrained TorchVision models for action recognition.
|
||||
|
||||
This class provides an interface for video classification using various pretrained models from TorchVision's video
|
||||
model collection, supporting models like S3D, R3D, Swin3D, and MViT architectures.
|
||||
|
||||
Attributes:
|
||||
model (torch.nn.Module): The loaded TorchVision model for video classification.
|
||||
weights (torchvision.models.video.Weights): The weights used for the model.
|
||||
device (torch.device): The device on which the model is loaded.
|
||||
|
||||
Methods:
|
||||
available_model_names: Returns a list of available model names.
|
||||
preprocess_crops_for_video_cls: Preprocesses crops for video classification.
|
||||
__call__: Performs inference on the given sequences.
|
||||
postprocess: Postprocesses the model's output.
|
||||
|
||||
Examples:
|
||||
>>> classifier = TorchVisionVideoClassifier("s3d", device="cpu")
|
||||
>>> crops = [np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8) for _ in range(8)]
|
||||
>>> tensor = classifier.preprocess_crops_for_video_cls(crops)
|
||||
>>> outputs = classifier(tensor)
|
||||
>>> labels, confidences = classifier.postprocess(outputs)
|
||||
|
||||
References:
|
||||
https://pytorch.org/vision/stable/
|
||||
"""
|
||||
|
||||
from torchvision.models.video import (
|
||||
MViT_V1_B_Weights,
|
||||
MViT_V2_S_Weights,
|
||||
R3D_18_Weights,
|
||||
S3D_Weights,
|
||||
Swin3D_B_Weights,
|
||||
Swin3D_T_Weights,
|
||||
mvit_v1_b,
|
||||
mvit_v2_s,
|
||||
r3d_18,
|
||||
s3d,
|
||||
swin3d_b,
|
||||
swin3d_t,
|
||||
)
|
||||
|
||||
model_name_to_model_and_weights = {
|
||||
"s3d": (s3d, S3D_Weights.DEFAULT),
|
||||
"r3d_18": (r3d_18, R3D_18_Weights.DEFAULT),
|
||||
"swin3d_t": (swin3d_t, Swin3D_T_Weights.DEFAULT),
|
||||
"swin3d_b": (swin3d_b, Swin3D_B_Weights.DEFAULT),
|
||||
"mvit_v1_b": (mvit_v1_b, MViT_V1_B_Weights.DEFAULT),
|
||||
"mvit_v2_s": (mvit_v2_s, MViT_V2_S_Weights.DEFAULT),
|
||||
}
|
||||
|
||||
def __init__(self, model_name: str, device: str | torch.device = ""):
|
||||
"""Initialize the VideoClassifier with the specified model name and device.
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model to use. Must be one of the available models.
|
||||
device (str | torch.device): The device to run the model on.
|
||||
"""
|
||||
if model_name not in self.model_name_to_model_and_weights:
|
||||
raise ValueError(f"Invalid model name '{model_name}'. Available models: {self.available_model_names()}")
|
||||
model, self.weights = self.model_name_to_model_and_weights[model_name]
|
||||
self.device = select_device(device)
|
||||
self.model = model(weights=self.weights).to(self.device).eval()
|
||||
|
||||
@staticmethod
|
||||
def available_model_names() -> list[str]:
|
||||
"""Get the list of available model names.
|
||||
|
||||
Returns:
|
||||
(list[str]): List of available model names that can be used with this classifier.
|
||||
"""
|
||||
return list(TorchVisionVideoClassifier.model_name_to_model_and_weights.keys())
|
||||
|
||||
def preprocess_crops_for_video_cls(
|
||||
self, crops: list[np.ndarray], input_size: list[int] | None = None
|
||||
) -> torch.Tensor:
|
||||
"""Preprocess a list of crops for video classification.
|
||||
|
||||
Args:
|
||||
crops (list[np.ndarray]): List of crops to preprocess. Each crop should have dimensions (H, W, C).
|
||||
input_size (list[int], optional): The target input size for the model.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): Preprocessed crops as a tensor with dimensions (1, T, C, H, W).
|
||||
"""
|
||||
if input_size is None:
|
||||
input_size = [224, 224]
|
||||
from torchvision.transforms import v2
|
||||
|
||||
transform = v2.Compose(
|
||||
[
|
||||
v2.ToDtype(torch.float32, scale=True),
|
||||
v2.Resize(input_size, antialias=True),
|
||||
v2.Normalize(mean=self.weights.transforms().mean, std=self.weights.transforms().std),
|
||||
]
|
||||
)
|
||||
|
||||
processed_crops = [transform(torch.from_numpy(crop).permute(2, 0, 1)) for crop in crops]
|
||||
return torch.stack(processed_crops).unsqueeze(0).permute(0, 2, 1, 3, 4).to(self.device)
|
||||
|
||||
def __call__(self, sequences: torch.Tensor) -> torch.Tensor:
|
||||
"""Perform inference on the given sequences.
|
||||
|
||||
Args:
|
||||
sequences (torch.Tensor): The input sequences for the model with dimensions (B, T, C, H, W) for batched
|
||||
video frames or (T, C, H, W) for single video frames.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): The model's output logits.
|
||||
"""
|
||||
with torch.inference_mode():
|
||||
return self.model(sequences)
|
||||
|
||||
def postprocess(self, outputs: torch.Tensor) -> tuple[list[str], list[float]]:
|
||||
"""Postprocess the model's batch output.
|
||||
|
||||
Args:
|
||||
outputs (torch.Tensor): The model's output logits.
|
||||
|
||||
Returns:
|
||||
pred_labels (list[str]): The predicted labels.
|
||||
pred_confs (list[float]): The predicted confidences.
|
||||
"""
|
||||
pred_labels = []
|
||||
pred_confs = []
|
||||
for output in outputs:
|
||||
pred_class = output.argmax(0).item()
|
||||
pred_label = self.weights.meta["categories"][pred_class]
|
||||
pred_labels.append(pred_label)
|
||||
pred_conf = output.softmax(0)[pred_class].item()
|
||||
pred_confs.append(pred_conf)
|
||||
|
||||
return pred_labels, pred_confs
|
||||
|
||||
|
||||
class HuggingFaceVideoClassifier:
|
||||
"""Zero-shot video classifier using Hugging Face transformer models.
|
||||
|
||||
This class provides an interface for zero-shot video classification using Hugging Face models, supporting custom
|
||||
label sets and various transformer architectures for video understanding.
|
||||
|
||||
Attributes:
|
||||
fp16 (bool): Whether to use FP16 for inference.
|
||||
labels (list[str]): List of labels for zero-shot classification.
|
||||
device (torch.device): The device on which the model is loaded.
|
||||
processor (transformers.AutoProcessor): The processor for the model.
|
||||
model (transformers.AutoModel): The loaded Hugging Face model.
|
||||
|
||||
Methods:
|
||||
preprocess_crops_for_video_cls: Preprocesses crops for video classification.
|
||||
__call__: Performs inference on the given sequences.
|
||||
postprocess: Postprocesses the model's output.
|
||||
|
||||
Examples:
|
||||
>>> labels = ["walking", "running", "dancing"]
|
||||
>>> classifier = HuggingFaceVideoClassifier(labels, device="cpu")
|
||||
>>> crops = [np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8) for _ in range(8)]
|
||||
>>> tensor = classifier.preprocess_crops_for_video_cls(crops)
|
||||
>>> outputs = classifier(tensor)
|
||||
>>> labels, confidences = classifier.postprocess(outputs)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
labels: list[str],
|
||||
model_name: str = "microsoft/xclip-base-patch16-zero-shot",
|
||||
device: str | torch.device = "",
|
||||
fp16: bool = False,
|
||||
):
|
||||
"""Initialize the HuggingFaceVideoClassifier with the specified model name.
|
||||
|
||||
Args:
|
||||
labels (list[str]): List of labels for zero-shot classification.
|
||||
model_name (str): The name of the model to use.
|
||||
device (str | torch.device): The device to run the model on.
|
||||
fp16 (bool): Whether to use FP16 for inference.
|
||||
"""
|
||||
self.fp16 = fp16
|
||||
self.labels = labels
|
||||
self.device = select_device(device)
|
||||
self.processor = AutoProcessor.from_pretrained(model_name)
|
||||
model = AutoModel.from_pretrained(model_name).to(self.device)
|
||||
if fp16:
|
||||
model = model.half()
|
||||
self.model = model.eval()
|
||||
|
||||
def preprocess_crops_for_video_cls(
|
||||
self, crops: list[np.ndarray], input_size: list[int] | None = None
|
||||
) -> torch.Tensor:
|
||||
"""Preprocess a list of crops for video classification.
|
||||
|
||||
Args:
|
||||
crops (list[np.ndarray]): List of crops to preprocess. Each crop should have dimensions (H, W, C).
|
||||
input_size (list[int], optional): The target input size for the model.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): Preprocessed crops as a tensor with dimensions (1, T, C, H, W).
|
||||
"""
|
||||
if input_size is None:
|
||||
input_size = [224, 224]
|
||||
from torchvision import transforms
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Lambda(lambda x: x.float() / 255.0),
|
||||
transforms.Resize(input_size),
|
||||
transforms.Normalize(
|
||||
mean=self.processor.image_processor.image_mean, std=self.processor.image_processor.image_std
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
processed_crops = [transform(torch.from_numpy(crop).permute(2, 0, 1)) for crop in crops] # (T, C, H, W)
|
||||
output = torch.stack(processed_crops).unsqueeze(0).to(self.device) # (1, T, C, H, W)
|
||||
if self.fp16:
|
||||
output = output.half()
|
||||
return output
|
||||
|
||||
def __call__(self, sequences: torch.Tensor) -> torch.Tensor:
|
||||
"""Perform inference on the given sequences.
|
||||
|
||||
Args:
|
||||
sequences (torch.Tensor): Batched input video frames with shape (B, T, H, W, C).
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): The model's output logits.
|
||||
"""
|
||||
input_ids = self.processor(text=self.labels, return_tensors="pt", padding=True)["input_ids"].to(self.device)
|
||||
|
||||
inputs = {"pixel_values": sequences, "input_ids": input_ids}
|
||||
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(**inputs)
|
||||
|
||||
return outputs.logits_per_video
|
||||
|
||||
def postprocess(self, outputs: torch.Tensor) -> tuple[list[list[str]], list[list[float]]]:
|
||||
"""Postprocess the model's batch output.
|
||||
|
||||
Args:
|
||||
outputs (torch.Tensor): The model's output logits.
|
||||
|
||||
Returns:
|
||||
pred_labels (list[list[str]]): The predicted top2 labels for each sample.
|
||||
pred_confs (list[list[float]]): The predicted top2 confidences for each sample.
|
||||
"""
|
||||
pred_labels = []
|
||||
pred_confs = []
|
||||
|
||||
with torch.no_grad():
|
||||
logits_per_video = outputs # Assuming outputs is already the logits tensor
|
||||
probs = logits_per_video.softmax(dim=-1) # Use softmax to convert logits to probabilities
|
||||
|
||||
for prob in probs:
|
||||
top2_indices = prob.topk(2).indices.tolist()
|
||||
top2_labels = [self.labels[idx] for idx in top2_indices]
|
||||
top2_confs = prob[top2_indices].tolist()
|
||||
pred_labels.append(top2_labels)
|
||||
pred_confs.append(top2_confs)
|
||||
|
||||
return pred_labels, pred_confs
|
||||
|
||||
|
||||
def crop_and_pad(frame: np.ndarray, box: list[float], margin_percent: int) -> np.ndarray:
|
||||
"""Crop box with margin and take square crop from frame.
|
||||
|
||||
Args:
|
||||
frame (np.ndarray): The input frame to crop from.
|
||||
box (list[float]): The bounding box coordinates [x1, y1, x2, y2].
|
||||
margin_percent (int): The percentage of margin to add around the box.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The cropped and resized square image.
|
||||
"""
|
||||
x1, y1, x2, y2 = map(int, box)
|
||||
w, h = x2 - x1, y2 - y1
|
||||
|
||||
# Add margin
|
||||
margin_x, margin_y = int(w * margin_percent / 100), int(h * margin_percent / 100)
|
||||
x1, y1 = max(0, x1 - margin_x), max(0, y1 - margin_y)
|
||||
x2, y2 = min(frame.shape[1], x2 + margin_x), min(frame.shape[0], y2 + margin_y)
|
||||
|
||||
# Take square crop from frame
|
||||
size = max(y2 - y1, x2 - x1)
|
||||
center_y, center_x = (y1 + y2) // 2, (x1 + x2) // 2
|
||||
half_size = size // 2
|
||||
square_crop = frame[
|
||||
max(0, center_y - half_size) : min(frame.shape[0], center_y + half_size),
|
||||
max(0, center_x - half_size) : min(frame.shape[1], center_x + half_size),
|
||||
]
|
||||
|
||||
return cv2.resize(square_crop, (224, 224), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
|
||||
def run(
|
||||
weights: str = "yolo11n.pt",
|
||||
device: str = "",
|
||||
source: str = "https://www.youtube.com/watch?v=dQw4w9WgXcQ",
|
||||
output_path: str | None = None,
|
||||
crop_margin_percentage: int = 10,
|
||||
num_video_sequence_samples: int = 8,
|
||||
skip_frame: int = 2,
|
||||
video_cls_overlap_ratio: float = 0.25,
|
||||
fp16: bool = False,
|
||||
video_classifier_model: str = "microsoft/xclip-base-patch32",
|
||||
labels: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Run action recognition on a video source using YOLO for object detection and a video classifier.
|
||||
|
||||
Args:
|
||||
weights (str): Path to the YOLO model weights.
|
||||
device (str): Device to run the model on. Use 'cuda' for NVIDIA GPU, 'mps' for Apple Silicon, or 'cpu'.
|
||||
source (str): Path to mp4 video file or YouTube URL.
|
||||
output_path (str, optional): Path to save the output video.
|
||||
crop_margin_percentage (int): Percentage of margin to add around detected objects.
|
||||
num_video_sequence_samples (int): Number of video frames to use for classification.
|
||||
skip_frame (int): Number of frames to skip between detections.
|
||||
video_cls_overlap_ratio (float): Overlap ratio between video sequences.
|
||||
fp16 (bool): Whether to use half-precision floating point.
|
||||
video_classifier_model (str): Name or path of the video classifier model.
|
||||
labels (list[str], optional): List of labels for zero-shot classification.
|
||||
"""
|
||||
if labels is None:
|
||||
labels = [
|
||||
"walking",
|
||||
"running",
|
||||
"brushing teeth",
|
||||
"looking into phone",
|
||||
"weight lifting",
|
||||
"cooking",
|
||||
"sitting",
|
||||
]
|
||||
# Initialize models and device
|
||||
device = select_device(device)
|
||||
yolo_model = YOLO(weights).to(device)
|
||||
if video_classifier_model in TorchVisionVideoClassifier.available_model_names():
|
||||
print("'fp16' is not supported for TorchVisionVideoClassifier. Setting fp16 to False.")
|
||||
print(
|
||||
"'labels' is not used for TorchVisionVideoClassifier. Ignoring the provided labels and using Kinetics-400 labels."
|
||||
)
|
||||
video_classifier = TorchVisionVideoClassifier(video_classifier_model, device=device)
|
||||
else:
|
||||
video_classifier = HuggingFaceVideoClassifier(
|
||||
labels, model_name=video_classifier_model, device=device, fp16=fp16
|
||||
)
|
||||
|
||||
# Initialize video capture
|
||||
if source.startswith("http") and urlparse(source).hostname in {"www.youtube.com", "youtube.com", "youtu.be"}:
|
||||
source = get_best_youtube_url(source)
|
||||
elif not source.endswith(".mp4"):
|
||||
raise ValueError("Invalid source. Supported sources are YouTube URLs and MP4 files.")
|
||||
cap = cv2.VideoCapture(source)
|
||||
|
||||
# Get video properties
|
||||
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
|
||||
# Initialize VideoWriter
|
||||
if output_path is not None:
|
||||
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
|
||||
|
||||
# Initialize track history
|
||||
track_history = defaultdict(list)
|
||||
frame_counter = 0
|
||||
|
||||
track_ids_to_infer = []
|
||||
crops_to_infer = []
|
||||
pred_labels = []
|
||||
pred_confs = []
|
||||
|
||||
while cap.isOpened():
|
||||
success, frame = cap.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_counter += 1
|
||||
|
||||
# Run YOLO tracking
|
||||
results = yolo_model.track(frame, persist=True, classes=[0]) # Track only person class
|
||||
|
||||
if results[0].boxes.is_track:
|
||||
boxes = results[0].boxes.xyxy.cpu().numpy()
|
||||
track_ids = results[0].boxes.id.cpu().numpy()
|
||||
|
||||
# Visualize prediction
|
||||
annotator = Annotator(frame, line_width=3, font_size=10, pil=False)
|
||||
|
||||
if frame_counter % skip_frame == 0:
|
||||
crops_to_infer = []
|
||||
track_ids_to_infer = []
|
||||
|
||||
for box, track_id in zip(boxes, track_ids):
|
||||
if frame_counter % skip_frame == 0:
|
||||
crop = crop_and_pad(frame, box, crop_margin_percentage)
|
||||
track_history[track_id].append(crop)
|
||||
|
||||
if len(track_history[track_id]) > num_video_sequence_samples:
|
||||
track_history[track_id].pop(0)
|
||||
|
||||
if len(track_history[track_id]) == num_video_sequence_samples and frame_counter % skip_frame == 0:
|
||||
start_time = time.time()
|
||||
crops = video_classifier.preprocess_crops_for_video_cls(track_history[track_id])
|
||||
end_time = time.time()
|
||||
preprocess_time = end_time - start_time
|
||||
print(f"video cls preprocess time: {preprocess_time:.4f} seconds")
|
||||
crops_to_infer.append(crops)
|
||||
track_ids_to_infer.append(track_id)
|
||||
|
||||
if crops_to_infer and (
|
||||
not pred_labels
|
||||
or frame_counter % int(num_video_sequence_samples * skip_frame * (1 - video_cls_overlap_ratio)) == 0
|
||||
):
|
||||
crops_batch = torch.cat(crops_to_infer, dim=0)
|
||||
|
||||
start_inference_time = time.time()
|
||||
output_batch = video_classifier(crops_batch)
|
||||
end_inference_time = time.time()
|
||||
inference_time = end_inference_time - start_inference_time
|
||||
print(f"video cls inference time: {inference_time:.4f} seconds")
|
||||
|
||||
pred_labels, pred_confs = video_classifier.postprocess(output_batch)
|
||||
|
||||
if track_ids_to_infer and crops_to_infer:
|
||||
for box, track_id, pred_label, pred_conf in zip(boxes, track_ids_to_infer, pred_labels, pred_confs):
|
||||
top2_preds = sorted(zip(pred_label, pred_conf), key=lambda x: x[1], reverse=True)
|
||||
label_text = " | ".join([f"{label} ({conf:.2f})" for label, conf in top2_preds])
|
||||
annotator.box_label(box, label_text, color=(0, 0, 255))
|
||||
|
||||
# Write the annotated frame to the output video
|
||||
if output_path is not None:
|
||||
out.write(frame)
|
||||
|
||||
# Display the annotated frame
|
||||
cv2.imshow("YOLOv8 Tracking with S3D Classification", frame)
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
if output_path is not None:
|
||||
out.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def parse_opt() -> argparse.Namespace:
|
||||
"""Parse command line arguments for action recognition pipeline."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--weights", type=str, default="yolo11n.pt", help="ultralytics detector model path")
|
||||
parser.add_argument("--device", default="", help='cuda device, i.e. 0 or 0,1,2,3 or cpu/mps, "" for auto-detection')
|
||||
parser.add_argument(
|
||||
"--source",
|
||||
type=str,
|
||||
default="https://www.youtube.com/watch?v=dQw4w9WgXcQ",
|
||||
help="video file path or youtube URL",
|
||||
)
|
||||
parser.add_argument("--output-path", type=str, default="output_video.mp4", help="output video file path")
|
||||
parser.add_argument(
|
||||
"--crop-margin-percentage", type=int, default=10, help="percentage of margin to add around detected objects"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-video-sequence-samples", type=int, default=8, help="number of video frames to use for classification"
|
||||
)
|
||||
parser.add_argument("--skip-frame", type=int, default=2, help="number of frames to skip between detections")
|
||||
parser.add_argument(
|
||||
"--video-cls-overlap-ratio", type=float, default=0.25, help="overlap ratio between video sequences"
|
||||
)
|
||||
parser.add_argument("--fp16", action="store_true", help="use FP16 for inference")
|
||||
parser.add_argument(
|
||||
"--video-classifier-model", type=str, default="microsoft/xclip-base-patch32", help="video classifier model name"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--labels",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=["dancing", "singing a song"],
|
||||
help="labels for zero-shot video classification",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(opt: argparse.Namespace) -> None:
|
||||
"""Run the action recognition pipeline with parsed command line arguments."""
|
||||
run(**vars(opt))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
opt = parse_opt()
|
||||
main(opt)
|
||||
@@ -0,0 +1,4 @@
|
||||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
ultralytics
|
||||
transformers
|
||||
@@ -0,0 +1,28 @@
|
||||
cmake_minimum_required(VERSION 3.5)
|
||||
|
||||
project(Yolov8CPPInference VERSION 0.1)
|
||||
|
||||
set(CMAKE_INCLUDE_CURRENT_DIR ON)
|
||||
|
||||
# CUDA
|
||||
set(CUDA_TOOLKIT_ROOT_DIR "/usr/local/cuda")
|
||||
find_package(CUDA 11 REQUIRED)
|
||||
|
||||
set(CMAKE_CUDA_STANDARD 11)
|
||||
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
|
||||
# !CUDA
|
||||
|
||||
# OpenCV
|
||||
find_package(OpenCV REQUIRED)
|
||||
include_directories(${OpenCV_INCLUDE_DIRS})
|
||||
# !OpenCV
|
||||
|
||||
set(PROJECT_SOURCES
|
||||
main.cpp
|
||||
|
||||
inference.h
|
||||
inference.cpp
|
||||
)
|
||||
|
||||
add_executable(Yolov8CPPInference ${PROJECT_SOURCES})
|
||||
target_link_libraries(Yolov8CPPInference ${OpenCV_LIBS})
|
||||
@@ -0,0 +1,85 @@
|
||||
# YOLOv8/YOLOv5 C++ Inference with OpenCV DNN
|
||||
|
||||
This example demonstrates how to perform inference using Ultralytics YOLOv8 and YOLOv5 models in C++ leveraging the [OpenCV DNN module](https://docs.opencv.org/4.x/d6/d0f/group__dnn.html).
|
||||
|
||||
## 🛠️ Usage
|
||||
|
||||
Follow these steps to set up and run the C++ inference example:
|
||||
|
||||
```bash
|
||||
# 1. Clone the Ultralytics repository
|
||||
git clone https://github.com/ultralytics/ultralytics
|
||||
cd ultralytics
|
||||
|
||||
# 2. Install Ultralytics Python package (needed for exporting models)
|
||||
pip install .
|
||||
|
||||
# 3. Navigate to the C++ example directory
|
||||
cd examples/YOLOv8-CPP-Inference
|
||||
|
||||
# 4. Export Models: Add yolov8*.onnx and/or yolov5*.onnx models (see export instructions below)
|
||||
# Place the exported ONNX models in the current directory (YOLOv8-CPP-Inference).
|
||||
|
||||
# 5. Update Source Code: Edit main.cpp and set the 'projectBasePath' variable
|
||||
# to the absolute path of the 'YOLOv8-CPP-Inference' directory on your system.
|
||||
# Example: std::string projectBasePath = "/path/to/your/ultralytics/examples/YOLOv8-CPP-Inference";
|
||||
|
||||
# 6. Configure OpenCV DNN Backend (Optional - CUDA):
|
||||
# - The default CMakeLists.txt attempts to use CUDA for GPU acceleration with OpenCV DNN.
|
||||
# - If your OpenCV build doesn't support CUDA/cuDNN, or you want CPU inference,
|
||||
# remove the CUDA-related lines from CMakeLists.txt.
|
||||
|
||||
# 7. Build the project
|
||||
mkdir build
|
||||
cd build
|
||||
cmake ..
|
||||
make
|
||||
|
||||
# 8. Run the inference executable
|
||||
./Yolov8CPPInference
|
||||
```
|
||||
|
||||
## ✨ Exporting YOLOv8 and YOLOv5 Models
|
||||
|
||||
You need to export your trained PyTorch models to the [ONNX](https://onnx.ai/) format to use them with OpenCV DNN.
|
||||
|
||||
**Exporting Ultralytics YOLOv8 Models:**
|
||||
|
||||
Use the Ultralytics CLI to export. Ensure you specify the desired `imgsz` and `opset`. For compatibility with this example, `opset=12` is recommended.
|
||||
|
||||
```bash
|
||||
yolo export model=yolov8s.pt imgsz=640,480 format=onnx opset=12 # Example: 640x480 resolution
|
||||
```
|
||||
|
||||
**Exporting YOLOv5 Models:**
|
||||
|
||||
Use the `export.py` script from the YOLOv5 repository structure (included within the cloned `ultralytics` repo).
|
||||
|
||||
```bash
|
||||
# Assuming you are in the 'ultralytics' base directory after cloning
|
||||
python export.py --weights yolov5s.pt --imgsz 640 480 --include onnx --opset 12 # Example: 640x480 resolution
|
||||
```
|
||||
|
||||
Place the generated `.onnx` files (e.g., `yolov8s.onnx`, `yolov5s.onnx`) into the `ultralytics/examples/YOLOv8-CPP-Inference/` directory.
|
||||
|
||||
**Example Output:**
|
||||
|
||||
_yolov8s.onnx:_
|
||||
|
||||
|
||||
|
||||
_yolov5s.onnx:_
|
||||
|
||||
|
||||
|
||||
## 📝 Notes
|
||||
|
||||
- This repository utilizes the [OpenCV DNN API](https://docs.opencv.org/4.x/d6/d0f/group__dnn.html) to run [ONNX](https://onnx.ai/) exported models of YOLOv5 and Ultralytics YOLOv8.
|
||||
- While not explicitly tested, it might theoretically work for other YOLO architectures like YOLOv6 and YOLOv7 if their ONNX export formats are compatible.
|
||||
- The example models are exported with a rectangular resolution (640x480), but the code should handle models exported with different resolutions. Consider using techniques like [letterboxing](https://docs.ultralytics.com/modes/predict/#letterbox) if your input images have different aspect ratios than the model's training resolution, especially for square `imgsz` exports.
|
||||
- The `main` branch version includes a simple GUI wrapper using [Qt](https://www.qt.io/). However, the core logic resides in the `Inference` class (`inference.h`, `inference.cpp`).
|
||||
- A key part of the `Inference` class demonstrates how to handle the output differences between YOLOv5 and YOLOv8 models, effectively transposing YOLOv8's output format to match the structure expected from YOLOv5 for consistent post-processing.
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! If you find any issues or have suggestions for improvement, please feel free to open an issue or submit a pull request. See our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for more details.
|
||||
@@ -0,0 +1,197 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include "inference.h"
|
||||
|
||||
Inference::Inference(const std::string &onnxModelPath, const cv::Size &modelInputShape, const std::string &classesTxtFile, const bool &runWithCuda)
|
||||
{
|
||||
modelPath = onnxModelPath;
|
||||
modelShape = modelInputShape;
|
||||
classesPath = classesTxtFile;
|
||||
cudaEnabled = runWithCuda;
|
||||
|
||||
loadOnnxNetwork();
|
||||
// loadClassesFromFile(); The classes are hard-coded for this example
|
||||
}
|
||||
|
||||
std::vector<Detection> Inference::runInference(const cv::Mat &input)
|
||||
{
|
||||
cv::Mat modelInput = input;
|
||||
int pad_x, pad_y;
|
||||
float scale;
|
||||
if (letterBoxForSquare && modelShape.width == modelShape.height)
|
||||
modelInput = formatToSquare(modelInput, &pad_x, &pad_y, &scale);
|
||||
|
||||
cv::Mat blob;
|
||||
cv::dnn::blobFromImage(modelInput, blob, 1.0/255.0, modelShape, cv::Scalar(), true, false);
|
||||
net.setInput(blob);
|
||||
|
||||
std::vector<cv::Mat> outputs;
|
||||
net.forward(outputs, net.getUnconnectedOutLayersNames());
|
||||
|
||||
int rows = outputs[0].size[1];
|
||||
int dimensions = outputs[0].size[2];
|
||||
|
||||
bool yolov8 = false;
|
||||
// yolov5 has an output of shape (batchSize, 25200, 85) (Num classes + box[x,y,w,h] + confidence[c])
|
||||
// yolov8 has an output of shape (batchSize, 84, 8400) (Num classes + box[x,y,w,h])
|
||||
if (dimensions > rows) // Check if the shape[2] is more than shape[1] (yolov8)
|
||||
{
|
||||
yolov8 = true;
|
||||
rows = outputs[0].size[2];
|
||||
dimensions = outputs[0].size[1];
|
||||
|
||||
outputs[0] = outputs[0].reshape(1, dimensions);
|
||||
cv::transpose(outputs[0], outputs[0]);
|
||||
}
|
||||
float *data = (float *)outputs[0].data;
|
||||
|
||||
std::vector<int> class_ids;
|
||||
std::vector<float> confidences;
|
||||
std::vector<cv::Rect> boxes;
|
||||
|
||||
for (int i = 0; i < rows; ++i)
|
||||
{
|
||||
if (yolov8)
|
||||
{
|
||||
float *classes_scores = data+4;
|
||||
|
||||
cv::Mat scores(1, classes.size(), CV_32FC1, classes_scores);
|
||||
cv::Point class_id;
|
||||
double maxClassScore;
|
||||
|
||||
minMaxLoc(scores, 0, &maxClassScore, 0, &class_id);
|
||||
|
||||
if (maxClassScore > modelScoreThreshold)
|
||||
{
|
||||
confidences.push_back(maxClassScore);
|
||||
class_ids.push_back(class_id.x);
|
||||
|
||||
float x = data[0];
|
||||
float y = data[1];
|
||||
float w = data[2];
|
||||
float h = data[3];
|
||||
|
||||
int left = int((x - 0.5 * w - pad_x) / scale);
|
||||
int top = int((y - 0.5 * h - pad_y) / scale);
|
||||
|
||||
int width = int(w / scale);
|
||||
int height = int(h / scale);
|
||||
|
||||
boxes.push_back(cv::Rect(left, top, width, height));
|
||||
}
|
||||
}
|
||||
else // yolov5
|
||||
{
|
||||
float confidence = data[4];
|
||||
|
||||
if (confidence >= modelConfidenceThreshold)
|
||||
{
|
||||
float *classes_scores = data+5;
|
||||
|
||||
cv::Mat scores(1, classes.size(), CV_32FC1, classes_scores);
|
||||
cv::Point class_id;
|
||||
double max_class_score;
|
||||
|
||||
minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
|
||||
|
||||
if (max_class_score > modelScoreThreshold)
|
||||
{
|
||||
confidences.push_back(confidence);
|
||||
class_ids.push_back(class_id.x);
|
||||
|
||||
float x = data[0];
|
||||
float y = data[1];
|
||||
float w = data[2];
|
||||
float h = data[3];
|
||||
|
||||
int left = int((x - 0.5 * w - pad_x) / scale);
|
||||
int top = int((y - 0.5 * h - pad_y) / scale);
|
||||
|
||||
int width = int(w / scale);
|
||||
int height = int(h / scale);
|
||||
|
||||
boxes.push_back(cv::Rect(left, top, width, height));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
data += dimensions;
|
||||
}
|
||||
|
||||
std::vector<int> nms_result;
|
||||
cv::dnn::NMSBoxes(boxes, confidences, modelScoreThreshold, modelNMSThreshold, nms_result);
|
||||
|
||||
std::vector<Detection> detections{};
|
||||
for (unsigned long i = 0; i < nms_result.size(); ++i)
|
||||
{
|
||||
int idx = nms_result[i];
|
||||
|
||||
Detection result;
|
||||
result.class_id = class_ids[idx];
|
||||
result.confidence = confidences[idx];
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_int_distribution<int> dis(100, 255);
|
||||
result.color = cv::Scalar(dis(gen),
|
||||
dis(gen),
|
||||
dis(gen));
|
||||
|
||||
result.className = classes[result.class_id];
|
||||
result.box = boxes[idx];
|
||||
|
||||
detections.push_back(result);
|
||||
}
|
||||
|
||||
return detections;
|
||||
}
|
||||
|
||||
void Inference::loadClassesFromFile()
|
||||
{
|
||||
std::ifstream inputFile(classesPath);
|
||||
if (inputFile.is_open())
|
||||
{
|
||||
std::string classLine;
|
||||
while (std::getline(inputFile, classLine))
|
||||
classes.push_back(classLine);
|
||||
inputFile.close();
|
||||
}
|
||||
}
|
||||
|
||||
void Inference::loadOnnxNetwork()
|
||||
{
|
||||
net = cv::dnn::readNetFromONNX(modelPath);
|
||||
if (cudaEnabled)
|
||||
{
|
||||
std::cout << "\nRunning on CUDA" << std::endl;
|
||||
net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
|
||||
net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "\nRunning on CPU" << std::endl;
|
||||
net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
|
||||
net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
|
||||
}
|
||||
}
|
||||
|
||||
cv::Mat Inference::formatToSquare(const cv::Mat &source, int *pad_x, int *pad_y, float *scale)
|
||||
{
|
||||
int col = source.cols;
|
||||
int row = source.rows;
|
||||
int m_inputWidth = modelShape.width;
|
||||
int m_inputHeight = modelShape.height;
|
||||
|
||||
*scale = std::min(m_inputWidth / (float)col, m_inputHeight / (float)row);
|
||||
int resized_w = col * *scale;
|
||||
int resized_h = row * *scale;
|
||||
*pad_x = (m_inputWidth - resized_w) / 2;
|
||||
*pad_y = (m_inputHeight - resized_h) / 2;
|
||||
|
||||
cv::Mat resized;
|
||||
cv::resize(source, resized, cv::Size(resized_w, resized_h));
|
||||
cv::Mat result = cv::Mat::zeros(m_inputHeight, m_inputWidth, source.type());
|
||||
resized.copyTo(result(cv::Rect(*pad_x, *pad_y, resized_w, resized_h)));
|
||||
resized.release();
|
||||
return result;
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#ifndef INFERENCE_H
|
||||
#define INFERENCE_H
|
||||
|
||||
// Cpp native
|
||||
#include <fstream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <random>
|
||||
|
||||
// OpenCV / DNN / Inference
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <opencv2/dnn.hpp>
|
||||
|
||||
struct Detection
|
||||
{
|
||||
int class_id{0};
|
||||
std::string className{};
|
||||
float confidence{0.0};
|
||||
cv::Scalar color{};
|
||||
cv::Rect box{};
|
||||
};
|
||||
|
||||
class Inference
|
||||
{
|
||||
public:
|
||||
Inference(const std::string &onnxModelPath, const cv::Size &modelInputShape = {640, 640}, const std::string &classesTxtFile = "", const bool &runWithCuda = true);
|
||||
std::vector<Detection> runInference(const cv::Mat &input);
|
||||
|
||||
private:
|
||||
void loadClassesFromFile();
|
||||
void loadOnnxNetwork();
|
||||
cv::Mat formatToSquare(const cv::Mat &source, int *pad_x, int *pad_y, float *scale);
|
||||
|
||||
std::string modelPath{};
|
||||
std::string classesPath{};
|
||||
bool cudaEnabled{};
|
||||
|
||||
std::vector<std::string> classes{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"};
|
||||
|
||||
cv::Size2f modelShape{};
|
||||
|
||||
float modelConfidenceThreshold {0.25};
|
||||
float modelScoreThreshold {0.45};
|
||||
float modelNMSThreshold {0.50};
|
||||
|
||||
bool letterBoxForSquare = true;
|
||||
|
||||
cv::dnn::Net net;
|
||||
};
|
||||
|
||||
#endif // INFERENCE_H
|
||||
@@ -0,0 +1,72 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <getopt.h>
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
#include "inference.h"
|
||||
|
||||
using namespace std;
|
||||
using namespace cv;
|
||||
|
||||
int main(int argc, char **argv)
|
||||
{
|
||||
std::string projectBasePath = "/home/user/ultralytics"; // Set your ultralytics base path
|
||||
|
||||
bool runOnGPU = true;
|
||||
|
||||
//
|
||||
// Pass in either:
|
||||
//
|
||||
// "yolov8s.onnx" or "yolov5s.onnx"
|
||||
//
|
||||
// To run Inference with yolov8/yolov5 (ONNX)
|
||||
//
|
||||
|
||||
// Note that in this example the classes are hard-coded and 'classes.txt' is a place holder.
|
||||
Inference inf(projectBasePath + "/yolov8s.onnx", cv::Size(640, 640), "classes.txt", runOnGPU);
|
||||
|
||||
std::vector<std::string> imageNames;
|
||||
imageNames.push_back(projectBasePath + "/ultralytics/assets/bus.jpg");
|
||||
imageNames.push_back(projectBasePath + "/ultralytics/assets/zidane.jpg");
|
||||
|
||||
for (int i = 0; i < imageNames.size(); ++i)
|
||||
{
|
||||
cv::Mat frame = cv::imread(imageNames[i]);
|
||||
|
||||
// Inference starts here...
|
||||
std::vector<Detection> output = inf.runInference(frame);
|
||||
|
||||
int detections = output.size();
|
||||
std::cout << "Number of detections:" << detections << std::endl;
|
||||
|
||||
for (int i = 0; i < detections; ++i)
|
||||
{
|
||||
Detection detection = output[i];
|
||||
|
||||
cv::Rect box = detection.box;
|
||||
cv::Scalar color = detection.color;
|
||||
|
||||
// Detection box
|
||||
cv::rectangle(frame, box, color, 2);
|
||||
|
||||
// Detection box text
|
||||
std::string classString = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4);
|
||||
cv::Size textSize = cv::getTextSize(classString, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
|
||||
cv::Rect textBox(box.x, box.y - 40, textSize.width + 10, textSize.height + 20);
|
||||
|
||||
cv::rectangle(frame, textBox, color, cv::FILLED);
|
||||
cv::putText(frame, classString, cv::Point(box.x + 5, box.y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
|
||||
}
|
||||
// Inference ends here...
|
||||
|
||||
// This is only for preview purposes
|
||||
float scale = 0.8;
|
||||
cv::resize(frame, frame, cv::Size(frame.cols*scale, frame.rows*scale));
|
||||
cv::imshow("Inference", frame);
|
||||
|
||||
cv::waitKey(-1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
|
||||
|
||||
project(yolov8_libtorch_example)
|
||||
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
set(CMAKE_CXX_EXTENSIONS OFF)
|
||||
|
||||
|
||||
# -------------- OpenCV --------------
|
||||
set(OpenCV_DIR "/path/to/opencv/lib/cmake/opencv4")
|
||||
find_package(OpenCV REQUIRED)
|
||||
|
||||
message(STATUS "OpenCV library status:")
|
||||
message(STATUS " config: ${OpenCV_DIR}")
|
||||
message(STATUS " version: ${OpenCV_VERSION}")
|
||||
message(STATUS " libraries: ${OpenCV_LIBS}")
|
||||
message(STATUS " include path: ${OpenCV_INCLUDE_DIRS}")
|
||||
|
||||
include_directories(${OpenCV_INCLUDE_DIRS})
|
||||
|
||||
# -------------- libtorch --------------
|
||||
list(APPEND CMAKE_PREFIX_PATH "/path/to/libtorch")
|
||||
set(Torch_DIR "/path/to/libtorch/share/cmake/Torch")
|
||||
|
||||
find_package(Torch REQUIRED)
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
|
||||
message("${TORCH_LIBRARIES}")
|
||||
message("${TORCH_INCLUDE_DIRS}")
|
||||
|
||||
# The following code block is suggested to be used on Windows.
|
||||
# According to https://github.com/pytorch/pytorch/issues/25457,
|
||||
# the DLLs need to be copied to avoid memory errors.
|
||||
# if (MSVC)
|
||||
# file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
|
||||
# add_custom_command(TARGET yolov8_libtorch_example
|
||||
# POST_BUILD
|
||||
# COMMAND ${CMAKE_COMMAND} -E copy_if_different
|
||||
# ${TORCH_DLLS}
|
||||
# $<TARGET_FILE_DIR:yolov8_libtorch_example>)
|
||||
# endif (MSVC)
|
||||
|
||||
include_directories(${TORCH_INCLUDE_DIRS})
|
||||
|
||||
add_executable(yolov8_libtorch_inference "${CMAKE_CURRENT_SOURCE_DIR}/main.cc")
|
||||
target_link_libraries(yolov8_libtorch_inference ${TORCH_LIBRARIES} ${OpenCV_LIBS})
|
||||
set_property(TARGET yolov8_libtorch_inference PROPERTY CXX_STANDARD 17)
|
||||
@@ -0,0 +1,74 @@
|
||||
# YOLOv8 LibTorch Inference C++
|
||||
|
||||
This example demonstrates how to perform inference using [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) models in C++ with the [LibTorch (PyTorch C++ API)](https://docs.pytorch.org/cppdocs/). This allows for deploying YOLOv8 models in C++ environments for efficient execution.
|
||||
|
||||
## ⚙️ Dependencies
|
||||
|
||||
Ensure you have the following dependencies installed before proceeding:
|
||||
|
||||
| Dependency | Version | Resource |
|
||||
| :----------- | :------- | :------------------------------------------- |
|
||||
| OpenCV | >=4.0.0 | [https://opencv.org/](https://opencv.org/) |
|
||||
| C++ Standard | >=17 | [https://isocpp.org/](https://isocpp.org/) |
|
||||
| CMake | >=3.18 | [https://cmake.org/](https://cmake.org/) |
|
||||
| Libtorch | >=1.12.1 | [https://pytorch.org/](https://pytorch.org/) |
|
||||
|
||||
You can download the required version of LibTorch from the official [PyTorch](https://pytorch.org/) website. Make sure to select the correct version corresponding to your system and CUDA version (if using GPU).
|
||||
|
||||
## 🚀 Usage
|
||||
|
||||
Follow these steps to run the C++ inference example:
|
||||
|
||||
1. **Clone the Ultralytics Repository:**
|
||||
Use [Git](https://git-scm.com/) to clone the repository containing the example code and necessary files.
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ultralytics/ultralytics
|
||||
```
|
||||
|
||||
2. **Install Ultralytics:**
|
||||
Navigate to the cloned directory and install the `ultralytics` package using [pip](https://pip.pypa.io/en/stable/). This step is necessary for exporting the model. Refer to the [Ultralytics Quickstart Guide](https://docs.ultralytics.com/quickstart/) for detailed installation instructions.
|
||||
|
||||
```bash
|
||||
cd ultralytics
|
||||
pip install .
|
||||
```
|
||||
|
||||
3. **Navigate to the Example Directory:**
|
||||
Change the directory to the C++ LibTorch inference example.
|
||||
|
||||
```bash
|
||||
cd examples/YOLOv8-LibTorch-CPP-Inference
|
||||
```
|
||||
|
||||
4. **Build the Project:**
|
||||
Create a build directory, use [CMake](https://cmake.org/) to configure the project, and then compile it using [Make](https://www.gnu.org/software/make/). You might need to specify the path to your LibTorch and OpenCV installations if they are not found automatically by CMake.
|
||||
|
||||
```bash
|
||||
mkdir build
|
||||
cd build
|
||||
cmake .. # Add -DCMAKE_PREFIX_PATH=/path/to/libtorch;/path/to/opencv if needed
|
||||
make
|
||||
```
|
||||
|
||||
5. **Run the Inference:**
|
||||
Execute the compiled binary. The application will load the exported YOLOv8 model and perform inference on a sample image (`zidane.jpg` included in the root `ultralytics` directory) or video.
|
||||
```bash
|
||||
./yolov8_libtorch_inference
|
||||
```
|
||||
|
||||
## ✨ Exporting Ultralytics YOLOv8
|
||||
|
||||
To use an Ultralytics YOLOv8 model with LibTorch, you first need to export it to the [TorchScript](https://docs.pytorch.org/docs/stable/jit.html) format. TorchScript is a way to create serializable and optimizable models from PyTorch code.
|
||||
|
||||
Use the `yolo` [Command Line Interface (CLI)](https://docs.ultralytics.com/usage/cli/) provided by the `ultralytics` package to export the model. For example, to export the `yolov8s.pt` model with an input image size of 640x640:
|
||||
|
||||
```bash
|
||||
yolo export model=yolov8s.pt imgsz=640 format=torchscript
|
||||
```
|
||||
|
||||
This command will generate a `yolov8s.torchscript` file in the model's directory. This file contains the serialized model that can be loaded and executed by the C++ application using LibTorch. For more details on exporting models to various formats, see the [Ultralytics Export documentation](https://docs.ultralytics.com/modes/export/).
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions to enhance this example or add new features are welcome! Please see the [Ultralytics Contributing Guide](https://docs.ultralytics.com/help/contributing/) for guidelines on how to contribute to the project. Thank you for helping make Ultralytics YOLO the best Vision AI tool!
|
||||
@@ -0,0 +1,260 @@
|
||||
#include <iostream>
|
||||
|
||||
#include <opencv2/core.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include <torch/torch.h>
|
||||
#include <torch/script.h>
|
||||
|
||||
using torch::indexing::Slice;
|
||||
using torch::indexing::None;
|
||||
|
||||
|
||||
float generate_scale(cv::Mat& image, const std::vector<int>& target_size) {
|
||||
int origin_w = image.cols;
|
||||
int origin_h = image.rows;
|
||||
|
||||
int target_h = target_size[0];
|
||||
int target_w = target_size[1];
|
||||
|
||||
float ratio_h = static_cast<float>(target_h) / static_cast<float>(origin_h);
|
||||
float ratio_w = static_cast<float>(target_w) / static_cast<float>(origin_w);
|
||||
float resize_scale = std::min(ratio_h, ratio_w);
|
||||
return resize_scale;
|
||||
}
|
||||
|
||||
|
||||
float letterbox(cv::Mat &input_image, cv::Mat &output_image, const std::vector<int> &target_size) {
|
||||
if (input_image.cols == target_size[1] && input_image.rows == target_size[0]) {
|
||||
if (input_image.data == output_image.data) {
|
||||
return 1.;
|
||||
} else {
|
||||
output_image = input_image.clone();
|
||||
return 1.;
|
||||
}
|
||||
}
|
||||
|
||||
float resize_scale = generate_scale(input_image, target_size);
|
||||
int new_shape_w = std::round(input_image.cols * resize_scale);
|
||||
int new_shape_h = std::round(input_image.rows * resize_scale);
|
||||
float padw = (target_size[1] - new_shape_w) / 2.;
|
||||
float padh = (target_size[0] - new_shape_h) / 2.;
|
||||
|
||||
int top = std::round(padh - 0.1);
|
||||
int bottom = std::round(padh + 0.1);
|
||||
int left = std::round(padw - 0.1);
|
||||
int right = std::round(padw + 0.1);
|
||||
|
||||
cv::resize(input_image, output_image,
|
||||
cv::Size(new_shape_w, new_shape_h),
|
||||
0, 0, cv::INTER_AREA);
|
||||
|
||||
cv::copyMakeBorder(output_image, output_image, top, bottom, left, right,
|
||||
cv::BORDER_CONSTANT, cv::Scalar(114., 114., 114));
|
||||
return resize_scale;
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor xyxy2xywh(const torch::Tensor& x) {
|
||||
auto y = torch::empty_like(x);
|
||||
y.index_put_({"...", 0}, (x.index({"...", 0}) + x.index({"...", 2})).div(2));
|
||||
y.index_put_({"...", 1}, (x.index({"...", 1}) + x.index({"...", 3})).div(2));
|
||||
y.index_put_({"...", 2}, x.index({"...", 2}) - x.index({"...", 0}));
|
||||
y.index_put_({"...", 3}, x.index({"...", 3}) - x.index({"...", 1}));
|
||||
return y;
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor xywh2xyxy(const torch::Tensor& x) {
|
||||
auto y = torch::empty_like(x);
|
||||
auto dw = x.index({"...", 2}).div(2);
|
||||
auto dh = x.index({"...", 3}).div(2);
|
||||
y.index_put_({"...", 0}, x.index({"...", 0}) - dw);
|
||||
y.index_put_({"...", 1}, x.index({"...", 1}) - dh);
|
||||
y.index_put_({"...", 2}, x.index({"...", 0}) + dw);
|
||||
y.index_put_({"...", 3}, x.index({"...", 1}) + dh);
|
||||
return y;
|
||||
}
|
||||
|
||||
|
||||
// Reference: https://github.com/pytorch/vision/blob/main/torchvision/csrc/ops/cpu/nms_kernel.cpp
|
||||
torch::Tensor nms(const torch::Tensor& bboxes, const torch::Tensor& scores, float iou_threshold) {
|
||||
if (bboxes.numel() == 0)
|
||||
return torch::empty({0}, bboxes.options().dtype(torch::kLong));
|
||||
|
||||
auto x1_t = bboxes.select(1, 0).contiguous();
|
||||
auto y1_t = bboxes.select(1, 1).contiguous();
|
||||
auto x2_t = bboxes.select(1, 2).contiguous();
|
||||
auto y2_t = bboxes.select(1, 3).contiguous();
|
||||
|
||||
torch::Tensor areas_t = (x2_t - x1_t) * (y2_t - y1_t);
|
||||
|
||||
auto order_t = std::get<1>(
|
||||
scores.sort(/*stable=*/true, /*dim=*/0, /* descending=*/true));
|
||||
|
||||
auto ndets = bboxes.size(0);
|
||||
torch::Tensor suppressed_t = torch::zeros({ndets}, bboxes.options().dtype(torch::kByte));
|
||||
torch::Tensor keep_t = torch::zeros({ndets}, bboxes.options().dtype(torch::kLong));
|
||||
|
||||
auto suppressed = suppressed_t.data_ptr<uint8_t>();
|
||||
auto keep = keep_t.data_ptr<int64_t>();
|
||||
auto order = order_t.data_ptr<int64_t>();
|
||||
auto x1 = x1_t.data_ptr<float>();
|
||||
auto y1 = y1_t.data_ptr<float>();
|
||||
auto x2 = x2_t.data_ptr<float>();
|
||||
auto y2 = y2_t.data_ptr<float>();
|
||||
auto areas = areas_t.data_ptr<float>();
|
||||
|
||||
int64_t num_to_keep = 0;
|
||||
|
||||
for (int64_t _i = 0; _i < ndets; _i++) {
|
||||
auto i = order[_i];
|
||||
if (suppressed[i] == 1)
|
||||
continue;
|
||||
keep[num_to_keep++] = i;
|
||||
auto ix1 = x1[i];
|
||||
auto iy1 = y1[i];
|
||||
auto ix2 = x2[i];
|
||||
auto iy2 = y2[i];
|
||||
auto iarea = areas[i];
|
||||
|
||||
for (int64_t _j = _i + 1; _j < ndets; _j++) {
|
||||
auto j = order[_j];
|
||||
if (suppressed[j] == 1)
|
||||
continue;
|
||||
auto xx1 = std::max(ix1, x1[j]);
|
||||
auto yy1 = std::max(iy1, y1[j]);
|
||||
auto xx2 = std::min(ix2, x2[j]);
|
||||
auto yy2 = std::min(iy2, y2[j]);
|
||||
|
||||
auto w = std::max(static_cast<float>(0), xx2 - xx1);
|
||||
auto h = std::max(static_cast<float>(0), yy2 - yy1);
|
||||
auto inter = w * h;
|
||||
auto ovr = inter / (iarea + areas[j] - inter);
|
||||
if (ovr > iou_threshold)
|
||||
suppressed[j] = 1;
|
||||
}
|
||||
}
|
||||
return keep_t.narrow(0, 0, num_to_keep);
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor non_max_suppression(torch::Tensor& prediction, float conf_thres = 0.25, float iou_thres = 0.45, int max_det = 300) {
|
||||
auto bs = prediction.size(0);
|
||||
auto nc = prediction.size(1) - 4;
|
||||
auto nm = prediction.size(1) - nc - 4;
|
||||
auto mi = 4 + nc;
|
||||
auto xc = prediction.index({Slice(), Slice(4, mi)}).amax(1) > conf_thres;
|
||||
|
||||
prediction = prediction.transpose(-1, -2);
|
||||
prediction.index_put_({"...", Slice({None, 4})}, xywh2xyxy(prediction.index({"...", Slice(None, 4)})));
|
||||
|
||||
std::vector<torch::Tensor> output;
|
||||
for (int i = 0; i < bs; i++) {
|
||||
output.push_back(torch::zeros({0, 6 + nm}, prediction.device()));
|
||||
}
|
||||
|
||||
for (int xi = 0; xi < prediction.size(0); xi++) {
|
||||
auto x = prediction[xi];
|
||||
x = x.index({xc[xi]});
|
||||
auto x_split = x.split({4, nc, nm}, 1);
|
||||
auto box = x_split[0], cls = x_split[1], mask = x_split[2];
|
||||
auto [conf, j] = cls.max(1, true);
|
||||
x = torch::cat({box, conf, j.toType(torch::kFloat), mask}, 1);
|
||||
x = x.index({conf.view(-1) > conf_thres});
|
||||
int n = x.size(0);
|
||||
if (!n) { continue; }
|
||||
|
||||
// NMS
|
||||
auto c = x.index({Slice(), Slice{5, 6}}) * 7680;
|
||||
auto boxes = x.index({Slice(), Slice(None, 4)}) + c;
|
||||
auto scores = x.index({Slice(), 4});
|
||||
auto i = nms(boxes, scores, iou_thres);
|
||||
i = i.index({Slice(None, max_det)});
|
||||
output[xi] = x.index({i});
|
||||
}
|
||||
|
||||
return torch::stack(output);
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor clip_boxes(torch::Tensor& boxes, const std::vector<int>& shape) {
|
||||
boxes.index_put_({"...", 0}, boxes.index({"...", 0}).clamp(0, shape[1]));
|
||||
boxes.index_put_({"...", 1}, boxes.index({"...", 1}).clamp(0, shape[0]));
|
||||
boxes.index_put_({"...", 2}, boxes.index({"...", 2}).clamp(0, shape[1]));
|
||||
boxes.index_put_({"...", 3}, boxes.index({"...", 3}).clamp(0, shape[0]));
|
||||
return boxes;
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor scale_boxes(const std::vector<int>& img1_shape, torch::Tensor& boxes, const std::vector<int>& img0_shape) {
|
||||
auto gain = (std::min)((float)img1_shape[0] / img0_shape[0], (float)img1_shape[1] / img0_shape[1]);
|
||||
auto pad0 = std::round((float)(img1_shape[1] - img0_shape[1] * gain) / 2. - 0.1);
|
||||
auto pad1 = std::round((float)(img1_shape[0] - img0_shape[0] * gain) / 2. - 0.1);
|
||||
|
||||
boxes.index_put_({"...", 0}, boxes.index({"...", 0}) - pad0);
|
||||
boxes.index_put_({"...", 2}, boxes.index({"...", 2}) - pad0);
|
||||
boxes.index_put_({"...", 1}, boxes.index({"...", 1}) - pad1);
|
||||
boxes.index_put_({"...", 3}, boxes.index({"...", 3}) - pad1);
|
||||
boxes.index_put_({"...", Slice(None, 4)}, boxes.index({"...", Slice(None, 4)}).div(gain));
|
||||
return boxes;
|
||||
}
|
||||
|
||||
|
||||
int main() {
|
||||
// Device
|
||||
torch::Device device(torch::cuda::is_available() ? torch::kCUDA :torch::kCPU);
|
||||
|
||||
// Note that in this example the classes are hard-coded
|
||||
std::vector<std::string> classes {"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant",
|
||||
"stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
|
||||
"giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite",
|
||||
"baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife",
|
||||
"spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
|
||||
"couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
|
||||
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"};
|
||||
|
||||
try {
|
||||
// Load the model (e.g. yolov8s.torchscript)
|
||||
std::string model_path = "/path/to/yolov8s.torchscript";
|
||||
torch::jit::script::Module yolo_model;
|
||||
yolo_model = torch::jit::load(model_path, device);
|
||||
yolo_model.eval();
|
||||
yolo_model.to(device, torch::kFloat32);
|
||||
|
||||
// Load image and preprocess
|
||||
cv::Mat image = cv::imread("/path/to/bus.jpg");
|
||||
cv::Mat input_image;
|
||||
letterbox(image, input_image, {640, 640});
|
||||
cv::cvtColor(input_image, input_image, cv::COLOR_BGR2RGB);
|
||||
|
||||
torch::Tensor image_tensor = torch::from_blob(input_image.data, {input_image.rows, input_image.cols, 3}, torch::kByte).to(device);
|
||||
image_tensor = image_tensor.toType(torch::kFloat32).div(255);
|
||||
image_tensor = image_tensor.permute({2, 0, 1});
|
||||
image_tensor = image_tensor.unsqueeze(0);
|
||||
std::vector<torch::jit::IValue> inputs {image_tensor};
|
||||
|
||||
// Inference
|
||||
torch::Tensor output = yolo_model.forward(inputs).toTensor().cpu();
|
||||
|
||||
// NMS
|
||||
auto keep = non_max_suppression(output)[0];
|
||||
auto boxes = keep.index({Slice(), Slice(None, 4)});
|
||||
keep.index_put_({Slice(), Slice(None, 4)}, scale_boxes({input_image.rows, input_image.cols}, boxes, {image.rows, image.cols}));
|
||||
|
||||
// Show the results
|
||||
for (int i = 0; i < keep.size(0); i++) {
|
||||
int x1 = keep[i][0].item().toFloat();
|
||||
int y1 = keep[i][1].item().toFloat();
|
||||
int x2 = keep[i][2].item().toFloat();
|
||||
int y2 = keep[i][3].item().toFloat();
|
||||
float conf = keep[i][4].item().toFloat();
|
||||
int cls = keep[i][5].item().toInt();
|
||||
std::cout << "Rect: [" << x1 << "," << y1 << "," << x2 << "," << y2 << "] Conf: " << conf << " Class: " << classes[cls] << std::endl;
|
||||
}
|
||||
} catch (const c10::Error& e) {
|
||||
std::cout << e.msg() << std::endl;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
cmake_minimum_required(VERSION 3.12)
|
||||
project(mnn_yolo_cpp)
|
||||
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")
|
||||
|
||||
include_directories(${CMAKE_CURRENT_LIST_DIR}/include/)
|
||||
|
||||
link_directories(${CMAKE_CURRENT_LIST_DIR}/libs)
|
||||
|
||||
add_executable("main" "${CMAKE_CURRENT_LIST_DIR}/main.cpp")
|
||||
add_executable("main_interpreter" "${CMAKE_CURRENT_LIST_DIR}/main_interpreter.cpp")
|
||||
|
||||
target_link_libraries("main" MNN MNN_Express MNNOpenCV)
|
||||
target_link_libraries("main_interpreter" MNN MNN_Express MNNOpenCV)
|
||||
185
algorithms/dms_yolo/code/examples/YOLOv8-MNN-CPP/README.md
Normal file
185
algorithms/dms_yolo/code/examples/YOLOv8-MNN-CPP/README.md
Normal file
@@ -0,0 +1,185 @@
|
||||
# YOLOv8 MNN Inference in C++
|
||||
|
||||
Welcome to the Ultralytics YOLOv8 MNN Inference example in C++! This guide will help you get started with leveraging the powerful [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) models using the [Alibaba MNN](https://mnn-docs.readthedocs.io/en/latest/) inference engine in your C++ projects. Whether you're looking to enhance performance on CPU hardware or add flexibility to your applications, this example provides a solid foundation. Learn more about optimizing models and deployment strategies on the [Ultralytics blog](https://www.ultralytics.com/blog).
|
||||
|
||||
## 🌟 Features
|
||||
|
||||
- 🚀 **Model Format Support**: Native support for the MNN format.
|
||||
- ⚡ **Precision Options**: Run models in **FP32**, **FP16** ([half-precision](https://www.ultralytics.com/glossary/half-precision)), and **INT8** ([model quantization](https://www.ultralytics.com/glossary/model-quantization)) precisions for optimized performance and reduced resource consumption.
|
||||
- 🔄 **Dynamic Shape Loading**: Easily handle models with dynamic input shapes, a common requirement in many [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) tasks.
|
||||
- 📦 **Flexible API Usage**: Choose between MNN's high-level [Express API](https://github.com/alibaba/MNN) for a user-friendly interface or the lower-level [Interpreter API](https://mnn-docs.readthedocs.io/en/latest/cpp/Interpreter.html) for fine-grained control.
|
||||
|
||||
## 📋 Dependencies
|
||||
|
||||
To ensure smooth execution, please make sure you have the following dependencies installed:
|
||||
|
||||
| Dependency | Version | Description |
|
||||
| :------------------------------------------------ | :------- | :------------------------------------------------------------------------------- |
|
||||
| [MNN](https://mnn-docs.readthedocs.io/en/latest/) | >=2.0.0 | The core inference engine from Alibaba. |
|
||||
| [C++](https://en.cppreference.com/w/) | >=14 | A modern C++ compiler supporting C++14 features. |
|
||||
| [CMake](https://cmake.org/documentation/) | >=3.12.0 | Cross-platform build system generator required for building MNN and the example. |
|
||||
| [OpenCV](https://opencv.org/) | Optional | Used for image loading and preprocessing within the example (built with MNN). |
|
||||
|
||||
## ⚙️ Build Instructions
|
||||
|
||||
Follow these steps to build the project:
|
||||
|
||||
1. Clone the Ultralytics repository:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ultralytics/ultralytics.git
|
||||
cd ultralytics/examples/YOLOv8-MNN-CPP
|
||||
```
|
||||
|
||||
2. Clone the [Alibaba MNN repository](https://github.com/alibaba/MNN):
|
||||
|
||||
```bash
|
||||
git clone https://github.com/alibaba/MNN.git
|
||||
cd MNN
|
||||
```
|
||||
|
||||
3. Build the MNN library:
|
||||
|
||||
```bash
|
||||
# Create build directory
|
||||
mkdir build && cd build
|
||||
|
||||
# Configure CMake (enable OpenCV integration, disable shared libs, enable image codecs)
|
||||
cmake -DMNN_BUILD_OPENCV=ON -DBUILD_SHARED_LIBS=OFF -DMNN_IMGCODECS=ON ..
|
||||
|
||||
# Build the library (use -j flag for parallel compilation)
|
||||
make -j$(nproc) # Use nproc for Linux, sysctl -n hw.ncpu for macOS
|
||||
```
|
||||
|
||||
**Note:** If you encounter issues during the build process, consult the official [MNN documentation](https://mnn-docs.readthedocs.io/en/latest/) for detailed build instructions and troubleshooting tips.
|
||||
|
||||
4. Copy the required MNN libraries and headers to the example project directory:
|
||||
|
||||
```bash
|
||||
# Navigate back to the example directory
|
||||
cd ../..
|
||||
|
||||
# Create directories for libraries and headers if they don't exist
|
||||
mkdir -p libs include
|
||||
|
||||
# Copy static libraries
|
||||
cp MNN/build/libMNN.a libs/ # Main MNN library
|
||||
cp MNN/build/express/libMNN_Express.a libs/ # MNN Express API library
|
||||
cp MNN/build/tools/cv/libMNNOpenCV.a libs/ # MNN OpenCV wrapper library
|
||||
|
||||
# Copy header files
|
||||
cp -r MNN/include .
|
||||
cp -r MNN/tools/cv/include . # MNN OpenCV wrapper headers
|
||||
```
|
||||
|
||||
**Note:**
|
||||
- The library file extensions (`.a` for static) and paths might vary based on your operating system (e.g., use `.lib` on Windows) and build configuration. Adjust the commands accordingly.
|
||||
- This example uses static linking (`.a` files). If you built shared libraries (`.so`, `.dylib`, `.dll`), ensure they are correctly placed or accessible in your system's library path.
|
||||
|
||||
5. Create a build directory for the example project and compile using CMake:
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
cmake ..
|
||||
make
|
||||
```
|
||||
|
||||
## 🔄 Exporting YOLOv8 Models
|
||||
|
||||
To use your Ultralytics YOLOv8 model with this C++ example, you first need to export it to the MNN format. This can be done easily using the `yolo export` command provided by the Ultralytics Python package.
|
||||
|
||||
Refer to the [Ultralytics Export documentation](https://docs.ultralytics.com/modes/export/) for detailed instructions and options.
|
||||
|
||||
```bash
|
||||
# Export a YOLOv8n model to MNN format with input size 640x640
|
||||
yolo export model=yolov8n.pt imgsz=640 format=mnn
|
||||
```
|
||||
|
||||
Alternatively, you can use the `MNNConvert` tool provided by MNN:
|
||||
|
||||
```bash
|
||||
# Assuming MNNConvert is built and in your PATH or MNN build directory
|
||||
# Convert an ONNX model (first export YOLOv8 to ONNX)
|
||||
yolo export model=yolov8n.pt format=onnx
|
||||
./MNN/build/MNNConvert -f ONNX --modelFile yolov8n.onnx --MNNModel yolov8n.mnn --bizCode biz
|
||||
```
|
||||
|
||||
For more details on model conversion using MNN tools, see the [MNN Convert documentation](https://mnn-docs.readthedocs.io/en/latest/tools/convert.html).
|
||||
|
||||
## 🛠️ Usage
|
||||
|
||||
### Ultralytics CLI in Python (for comparison)
|
||||
|
||||
You can verify the exported MNN model using the Ultralytics Python package for a quick check.
|
||||
|
||||
Download an example image:
|
||||
|
||||
```bash
|
||||
wget https://ultralytics.com/images/bus.jpg
|
||||
```
|
||||
|
||||
Run prediction using the MNN model:
|
||||
|
||||
```bash
|
||||
yolo predict model=yolov8n.mnn source=bus.jpg
|
||||
```
|
||||
|
||||
Expected Python Output:
|
||||
|
||||
```
|
||||
ultralytics/examples/YOLOv8-MNN-CPP/assets/bus.jpg: 640x640 4 persons, 1 bus, 84.6ms
|
||||
Speed: 9.7ms preprocess, 128.7ms inference, 12.4ms postprocess per image at shape (1, 3, 640, 640)
|
||||
Results saved to runs/detect/predict
|
||||
```
|
||||
|
||||
_(Note: Speed and specific detections might vary based on hardware and model version)_
|
||||
|
||||
### MNN Express API in C++
|
||||
|
||||
This example uses the higher-level Express API for simpler inference code.
|
||||
|
||||
```bash
|
||||
./build/main yolov8n.mnn bus.jpg
|
||||
```
|
||||
|
||||
Expected C++ Express API Output:
|
||||
|
||||
```
|
||||
The device supports: i8sdot:0, fp16:0, i8mm: 0, sve2: 0, sme2: 0
|
||||
Detection: box = {48.63, 399.30, 243.65, 902.90}, class = person, score = 0.86
|
||||
Detection: box = {22.14, 228.36, 796.07, 749.74}, class = bus, score = 0.86
|
||||
Detection: box = {669.92, 375.82, 809.86, 874.41}, class = person, score = 0.86
|
||||
Detection: box = {216.01, 405.24, 346.36, 858.19}, class = person, score = 0.82
|
||||
Detection: box = {-0.11, 549.41, 62.05, 874.88}, class = person, score = 0.33
|
||||
Result image write to `mnn_yolov8_cpp.jpg`.
|
||||
Speed: 35.6ms preprocess, 386.0ms inference, 68.3ms postprocess
|
||||
```
|
||||
|
||||
_(Note: Speed and specific detections might vary based on hardware and MNN configuration)_
|
||||
|
||||
### MNN Interpreter API in C++
|
||||
|
||||
This example uses the lower-level Interpreter API, offering more control over the inference process.
|
||||
|
||||
```bash
|
||||
./build/main_interpreter yolov8n.mnn bus.jpg
|
||||
```
|
||||
|
||||
Expected C++ Interpreter API Output:
|
||||
|
||||
```
|
||||
The device supports: i8sdot:0, fp16:0, i8mm: 0, sve2: 0, sme2: 0
|
||||
Detection: box = {48.63, 399.30, 243.65, 902.90}, class = person, score = 0.86
|
||||
Detection: box = {22.14, 228.36, 796.07, 749.74}, class = bus, score = 0.86
|
||||
Detection: box = {669.92, 375.82, 809.86, 874.41}, class = person, score = 0.86
|
||||
Detection: box = {216.01, 405.24, 346.36, 858.19}, class = person, score = 0.82
|
||||
Result image written to `mnn_yolov8_cpp.jpg`.
|
||||
Speed: 26.0ms preprocess, 190.9ms inference, 58.9ms postprocess
|
||||
```
|
||||
|
||||
_(Note: Speed and specific detections might vary based on hardware and MNN configuration)_
|
||||
|
||||
## ❤️ Contributions
|
||||
|
||||
We hope this example helps you integrate Ultralytics YOLOv8 with MNN into your C++ projects effortlessly! Contributions to improve this example or add new features are highly welcome. Please see the [Ultralytics contribution guidelines](https://docs.ultralytics.com/help/contributing/) for more information on how to get involved.
|
||||
|
||||
For further guides, tutorials, and documentation on Ultralytics YOLO models and tools, visit the main [Ultralytics documentation](https://docs.ultralytics.com/). Happy coding! 🚀
|
||||
246
algorithms/dms_yolo/code/examples/YOLOv8-MNN-CPP/main.cpp
Normal file
246
algorithms/dms_yolo/code/examples/YOLOv8-MNN-CPP/main.cpp
Normal file
@@ -0,0 +1,246 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <algorithm>
|
||||
#include <regex>
|
||||
#include <sstream>
|
||||
#include <MNN/ImageProcess.hpp>
|
||||
#include <MNN/expr/Module.hpp>
|
||||
#include <MNN/expr/Executor.hpp>
|
||||
#include <MNN/expr/ExprCreator.hpp>
|
||||
#include <cv/cv.hpp>
|
||||
|
||||
using namespace MNN;
|
||||
using namespace MNN::Express;
|
||||
using namespace MNN::CV;
|
||||
|
||||
class Inference {
|
||||
public:
|
||||
|
||||
// Load model: Create runtime, set cache if needed, and load the model file.
|
||||
bool loadModel(const std::string &modelPath,
|
||||
int forwardType = MNN_FORWARD_CPU,
|
||||
int precision = 0,
|
||||
int thread = 4) {
|
||||
MNN::ScheduleConfig sConfig;
|
||||
sConfig.type = static_cast<MNNForwardType>(forwardType);
|
||||
sConfig.numThread = thread;
|
||||
BackendConfig bConfig;
|
||||
bConfig.precision = static_cast<BackendConfig::PrecisionMode>(precision);
|
||||
sConfig.backendConfig = &bConfig;
|
||||
|
||||
std::shared_ptr<Executor::RuntimeManager> rtmgr(
|
||||
Executor::RuntimeManager::createRuntimeManager(sConfig)
|
||||
);
|
||||
if (rtmgr == nullptr) {
|
||||
MNN_ERROR("Empty RuntimeManager\n");
|
||||
return false;
|
||||
}
|
||||
rtmgr->setCache(".cachefile");
|
||||
net = std::shared_ptr<Module>(Module::load(std::vector<std::string>{},
|
||||
std::vector<std::string>{}, modelPath.c_str(), rtmgr));
|
||||
if (net == nullptr) {
|
||||
return false;
|
||||
}
|
||||
runtimeManager = rtmgr;
|
||||
const Module::Info* info = net->getInfo();
|
||||
if (info == nullptr) {
|
||||
MNN_ERROR("Empty Module Info\n");
|
||||
return false;
|
||||
}
|
||||
// Parse bizCode to extract class names.
|
||||
if (info->bizCode.empty()) {
|
||||
MNN_ERROR("Empty bizCode\n");
|
||||
classNames.clear();
|
||||
return false;
|
||||
}
|
||||
// Get imgsz from bizCode.
|
||||
auto imgsz_start = info->bizCode.find("\"imgsz\": [");
|
||||
if (imgsz_start == std::string::npos) {
|
||||
MNN_PRINT("No imgsz found in bizCode, setting classNames empty.\n");
|
||||
} else {
|
||||
auto imgsz_end = info->bizCode.find("]", imgsz_start);
|
||||
if (imgsz_end == std::string::npos) {
|
||||
MNN_PRINT("No closing bracket for imgsz in bizCode, setting classNames empty.\n");
|
||||
} else {
|
||||
std::string imgszText = info->bizCode.substr(imgsz_start + 10, imgsz_end - imgsz_start - 10);
|
||||
std::vector<std::string> imgszVec;
|
||||
std::stringstream ss(imgszText);
|
||||
std::string item;
|
||||
while (std::getline(ss, item, ',')) {
|
||||
imgszVec.push_back(item);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Get names from bizCode.
|
||||
auto names_start = info->bizCode.find("\"names\": {");
|
||||
if (names_start == std::string::npos) {
|
||||
MNN_PRINT("No names found in bizCode, setting classNames empty.\n");
|
||||
classNames.clear();
|
||||
} else {
|
||||
auto names_end = info->bizCode.find("}", names_start);
|
||||
if (names_end == std::string::npos) {
|
||||
MNN_PRINT("No closing brace for names in bizCode, setting classNames empty.\n");
|
||||
classNames.clear();
|
||||
} else {
|
||||
std::string namesDict = info->bizCode.substr(names_start + 10, names_end - names_start - 10);
|
||||
parseClassNamesFromBizCode(namesDict);
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void parseImgszFromBizCode(const std::string& bizText) {
|
||||
std::regex rgx("\"imgsz\":\\s*\\[(\\d+),\\s*(\\d+)\\]");
|
||||
std::smatch match;
|
||||
if (std::regex_search(bizText, match, rgx)) {
|
||||
int ih = std::stoi(match[1].str());
|
||||
int iw = std::stoi(match[2].str());
|
||||
MNN_PRINT("Input size: %d x %d\n", iw, ih);
|
||||
} else {
|
||||
MNN_PRINT("No imgsz found in bizCode.\n");
|
||||
}
|
||||
}
|
||||
|
||||
void parseClassNamesFromBizCode(const std::string& bizText) {
|
||||
std::regex rgx("\"(\\d+)\"\\s*:\\s*\"([^\"]+)\"");
|
||||
std::smatch match;
|
||||
std::string s = bizText;
|
||||
classNames.clear();
|
||||
while (std::regex_search(s, match, rgx)) {
|
||||
int index = std::stoi(match[1].str());
|
||||
std::string name = match[2].str();
|
||||
if (classNames.size() <= static_cast<size_t>(index)) {
|
||||
classNames.resize(index + 1);
|
||||
}
|
||||
classNames[index] = name;
|
||||
s = match.suffix().str();
|
||||
}
|
||||
}
|
||||
|
||||
VARP preprocess(VARP &originalImage, float &scale) {
|
||||
auto dims = originalImage->getInfo()->dim;
|
||||
int ih = dims[0];
|
||||
int iw = dims[1];
|
||||
int targetWidth = std::stoi(imgszVec[0]);
|
||||
int targetHeight = std::stoi(imgszVec[1]);
|
||||
int len = ih > iw ? ih : iw;
|
||||
scale = static_cast<float>(len) / std::max(targetWidth, targetHeight);
|
||||
std::vector<int> padvals { 0, len - ih, 0, len - iw, 0, 0 };
|
||||
auto pads = _Const(static_cast<void*>(padvals.data()), {3, 2}, NCHW, halide_type_of<int>());
|
||||
auto image = _Pad(originalImage, pads, CONSTANT);
|
||||
image = resize(image, Size(targetWidth, targetHeight), 0, 0, INTER_LINEAR, -1, {0., 0., 0.}, {1./255., 1./255., 1./255.});
|
||||
auto input = _Unsqueeze(image, {0});
|
||||
input = _Convert(input, NC4HW4);
|
||||
return input;
|
||||
}
|
||||
|
||||
void runInference(VARP input) {
|
||||
std::vector<VARP> outputs = net->onForward({input});
|
||||
mOutput = outputs[0];
|
||||
}
|
||||
|
||||
void postprocess(float scale, VARP originalImage, float iouThreshold = 0.45, float scoreThreshold = 0.25) {
|
||||
auto output = _Convert(mOutput, NCHW);
|
||||
output = _Squeeze(output);
|
||||
// Expected output shape: [84, 8400]
|
||||
auto cx = _Gather(output, _Scalar<int>(0));
|
||||
auto cy = _Gather(output, _Scalar<int>(1));
|
||||
auto w = _Gather(output, _Scalar<int>(2));
|
||||
auto h = _Gather(output, _Scalar<int>(3));
|
||||
|
||||
std::vector<int> startvals { 4, 0 };
|
||||
auto start = _Const(static_cast<void*>(startvals.data()), {2}, NCHW, halide_type_of<int>());
|
||||
std::vector<int> sizevals { -1, -1 };
|
||||
auto size = _Const(static_cast<void*>(sizevals.data()), {2}, NCHW, halide_type_of<int>());
|
||||
auto probs = _Slice(output, start, size);
|
||||
// [cx, cy, w, h] -> [x1, y1, x2, y2]
|
||||
auto x1 = cx - w * _Const(0.5);
|
||||
auto y1 = cy - h * _Const(0.5);
|
||||
auto x2 = cx + w * _Const(0.5);
|
||||
auto y2 = cy + h * _Const(0.5);
|
||||
auto boxes = _Stack({x1, y1, x2, y2}, 1);
|
||||
auto scores = _ReduceMax(probs, {0});
|
||||
auto ids = _ArgMax(probs, 0);
|
||||
auto result_ids = _Nms(boxes, scores, 100, 0.45, 0.25);
|
||||
auto result_ptr = result_ids->readMap<int>();
|
||||
auto box_ptr = boxes->readMap<float>();
|
||||
auto ids_ptr = ids->readMap<int>();
|
||||
auto score_ptr = scores->readMap<float>();
|
||||
for (int i = 0; i < 100; i++) {
|
||||
auto idx = result_ptr[i];
|
||||
if (idx < 0) break;
|
||||
auto x1 = box_ptr[idx * 4 + 0] * scale;
|
||||
auto y1 = box_ptr[idx * 4 + 1] * scale;
|
||||
auto x2 = box_ptr[idx * 4 + 2] * scale;
|
||||
auto y2 = box_ptr[idx * 4 + 3] * scale;
|
||||
auto class_idx = ids_ptr[idx];
|
||||
auto score = score_ptr[idx];
|
||||
printf("Detection: box = {%.2f, %.2f, %.2f, %.2f}, class = %s, score = %.2f\n",
|
||||
x1, y1, x2, y2, classNames[class_idx].c_str(), score);
|
||||
rectangle(originalImage, { x1, y1 }, { x2, y2 }, { 0, 255, 0 }, 2);
|
||||
}
|
||||
if (imwrite("mnn_yolov8_cpp.jpg", originalImage)) {
|
||||
MNN_PRINT("Result image written to `mnn_yolov8_cpp.jpg`.\n");
|
||||
}
|
||||
}
|
||||
|
||||
// Update runtime cache.
|
||||
void updateCache() {
|
||||
if (runtimeManager)
|
||||
runtimeManager->updateCache();
|
||||
}
|
||||
|
||||
private:
|
||||
std::shared_ptr<Module> net;
|
||||
VARP mOutput;
|
||||
std::shared_ptr<Executor::RuntimeManager> runtimeManager;
|
||||
std::vector<std::string> classNames;
|
||||
std::vector<std::string> imgszVec { "640", "640" };
|
||||
};
|
||||
|
||||
int main(int argc, const char* argv[]) {
|
||||
if (argc < 3) {
|
||||
MNN_PRINT("Usage: ./main yolov8n.mnn bus.jpg [forwardType] [precision] [thread]\n");
|
||||
return 0;
|
||||
}
|
||||
int thread = 4;
|
||||
int precision = 0;
|
||||
int forwardType = MNN_FORWARD_CPU;
|
||||
if (argc >= 4) {
|
||||
forwardType = atoi(argv[3]);
|
||||
}
|
||||
if (argc >= 5) {
|
||||
precision = atoi(argv[4]);
|
||||
}
|
||||
if (argc >= 6) {
|
||||
thread = atoi(argv[5]);
|
||||
}
|
||||
|
||||
Inference infer;
|
||||
if (!infer.loadModel(argv[1], forwardType, precision, thread))
|
||||
return 1;
|
||||
|
||||
const clock_t t0 = clock();
|
||||
float scale = 1.0f;
|
||||
VARP originalImage = imread(argv[2]);
|
||||
VARP input = infer.preprocess(originalImage, scale);
|
||||
double preprocess_time = 1000.0 * (clock() - t0) / CLOCKS_PER_SEC;
|
||||
|
||||
const clock_t t1 = clock();
|
||||
infer.runInference(input);
|
||||
double inference_time = 1000.0 * (clock() - t1) / CLOCKS_PER_SEC;
|
||||
|
||||
const clock_t t2 = clock();
|
||||
infer.postprocess(scale, originalImage);
|
||||
double postprocess_time = 1000.0 * (clock() - t2) / CLOCKS_PER_SEC;
|
||||
|
||||
printf("Speed: %.1fms preprocess, %.1fms inference, %.1fms postprocess\n",
|
||||
preprocess_time, inference_time, postprocess_time);
|
||||
|
||||
infer.updateCache();
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <regex>
|
||||
#include <MNN/ImageProcess.hpp>
|
||||
#include <MNN/expr/Module.hpp>
|
||||
#include <MNN/expr/Executor.hpp>
|
||||
#include <MNN/expr/ExprCreator.hpp>
|
||||
#include <cv/cv.hpp>
|
||||
|
||||
using namespace MNN;
|
||||
using namespace MNN::Express;
|
||||
using namespace MNN::CV;
|
||||
|
||||
class Inference {
|
||||
public:
|
||||
Inference() : interpreter(nullptr), session(nullptr), inputTensor(nullptr) {
|
||||
inputDims = {1, 3, 640, 640};
|
||||
}
|
||||
|
||||
~Inference() {
|
||||
if(interpreter) {
|
||||
delete interpreter;
|
||||
interpreter = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
// Load model, create session, and resize the input tensor.
|
||||
bool loadModel(const std::string &modelPath,
|
||||
int forwardType = MNN_FORWARD_CPU,
|
||||
int precision = 1,
|
||||
int thread = 4) {
|
||||
MNN::ScheduleConfig sConfig;
|
||||
sConfig.type = static_cast<MNNForwardType>(forwardType);
|
||||
sConfig.numThread = thread;
|
||||
BackendConfig bConfig;
|
||||
bConfig.precision = static_cast<BackendConfig::PrecisionMode>(precision);
|
||||
sConfig.backendConfig = &bConfig;
|
||||
|
||||
interpreter = MNN::Interpreter::createFromFile(modelPath.c_str());
|
||||
if (!interpreter) {
|
||||
MNN_PRINT("Error: Failed to create interpreter from model file.\n");
|
||||
return false;
|
||||
}
|
||||
session = interpreter->createSession(sConfig);
|
||||
if(!session) {
|
||||
MNN_PRINT("Error: Failed to create session.\n");
|
||||
return false;
|
||||
}
|
||||
inputTensor = interpreter->getSessionInput(session, "images");
|
||||
interpreter->resizeTensor(inputTensor, inputDims);
|
||||
interpreter->resizeSession(session);
|
||||
std::string bizCode = interpreter->bizCode();
|
||||
|
||||
// Get names from bizCode.
|
||||
auto names_start = bizCode.find("\"names\": {");
|
||||
if (names_start == std::string::npos) {
|
||||
MNN_PRINT("No names found in bizCode, setting classNames empty.\n");
|
||||
classNames.clear();
|
||||
} else {
|
||||
auto names_end = bizCode.find("}", names_start);
|
||||
if (names_end == std::string::npos) {
|
||||
MNN_PRINT("No closing brace for names in bizCode, setting classNames empty.\n");
|
||||
classNames.clear();
|
||||
} else {
|
||||
std::string namesDict = bizCode.substr(names_start + 10, names_end - names_start - 10);
|
||||
parseClassNamesFromBizCode(namesDict);
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void parseClassNamesFromBizCode(const std::string& bizText) {
|
||||
std::regex rgx("\"(\\d+)\"\\s*:\\s*\"([^\"]+)\"");
|
||||
std::smatch match;
|
||||
std::string s = bizText;
|
||||
classNames.clear();
|
||||
while (std::regex_search(s, match, rgx)) {
|
||||
int index = std::stoi(match[1].str());
|
||||
std::string name = match[2].str();
|
||||
if (classNames.size() <= static_cast<size_t>(index)) {
|
||||
classNames.resize(index + 1);
|
||||
}
|
||||
classNames[index] = name;
|
||||
s = match.suffix().str();
|
||||
}
|
||||
}
|
||||
|
||||
VARP preprocess(VARP &originalImage, int targetSize, float &scale) {
|
||||
const auto dims = originalImage->getInfo()->dim;
|
||||
const int ih = dims[0], iw = dims[1];
|
||||
const int len = (ih >= iw ? ih : iw);
|
||||
scale = static_cast<float>(len) / targetSize;
|
||||
|
||||
// Use fixed-size array for padding values.
|
||||
int padvals[6] = { 0, len - ih, 0, len - iw, 0, 0 };
|
||||
auto pads = _Const(static_cast<void*>(padvals), {3, 2}, NCHW, halide_type_of<int>());
|
||||
auto padded = _Pad(originalImage, pads, CONSTANT);
|
||||
|
||||
auto resized = MNN::CV::resize(padded, MNN::CV::Size(targetSize, targetSize),
|
||||
0, 0, MNN::CV::INTER_LINEAR, -1,
|
||||
{0.f, 0.f, 0.f},
|
||||
{1.f/255, 1.f/255, 1.f/255});
|
||||
|
||||
// Chain unsqueeze and conversion
|
||||
auto input = _Unsqueeze(resized, {0});
|
||||
input = _Convert(input, NCHW);
|
||||
return input;
|
||||
}
|
||||
|
||||
// Run inference by copying preprocessed data into input tensor.
|
||||
void runInference(VARP input) {
|
||||
auto tmp_input = MNN::Tensor::create(inputDims, halide_type_of<float>(),
|
||||
const_cast<void*>(input->readMap<void>()),
|
||||
MNN::Tensor::CAFFE);
|
||||
inputTensor->copyFromHostTensor(tmp_input);
|
||||
interpreter->runSession(session);
|
||||
}
|
||||
|
||||
// Postprocess the output, perform NMS, and draw bounding boxes on originalImage.
|
||||
void postprocess(float scale, VARP originalImage, float modelScoreThreshold = 0.25, float modelNMSThreshold = 0.45) {
|
||||
auto outputTensor = interpreter->getSessionOutput(session, "output0");
|
||||
|
||||
// ---------------- Post Processing ----------------
|
||||
auto outputs = outputTensor->host<float>();
|
||||
auto outputVar = _Const(outputs, outputTensor->shape(), NCHW, halide_type_of<float>());
|
||||
auto output = _Squeeze(_Convert(outputVar, NCHW));
|
||||
|
||||
// Expected output shape: [84, 8400] where first 4 rows are [cx, cy, w, h].
|
||||
auto cx = _Gather(output, _Scalar<int>(0));
|
||||
auto cy = _Gather(output, _Scalar<int>(1));
|
||||
auto w = _Gather(output, _Scalar<int>(2));
|
||||
auto h = _Gather(output, _Scalar<int>(3));
|
||||
|
||||
// Slice probability values (starting at row 4).
|
||||
const int startArr[2] = { 4, 0 };
|
||||
const int sizeArr[2] = { -1, -1 };
|
||||
auto start = _Const(static_cast<void*>(const_cast<int*>(startArr)), {2}, NCHW, halide_type_of<int>());
|
||||
auto size = _Const(static_cast<void*>(const_cast<int*>(sizeArr)), {2}, NCHW, halide_type_of<int>());
|
||||
auto probs = _Slice(output, start, size);
|
||||
|
||||
// Convert [cx, cy, w, h] to [y1, x1, y2, x2] using half-width/height.
|
||||
auto half = _Const(0.5);
|
||||
auto x1 = cx - w * half;
|
||||
auto y1 = cy - h * half;
|
||||
auto x2 = cx + w * half;
|
||||
auto y2 = cy + h * half;
|
||||
auto boxes = _Stack({x1, y1, x2, y2}, 1);
|
||||
|
||||
auto scores = _ReduceMax(probs, {0});
|
||||
auto ids = _ArgMax(probs, 0);
|
||||
auto result_ids = _Nms(boxes, scores, 100, modelScoreThreshold, modelNMSThreshold);
|
||||
|
||||
auto result_ptr = result_ids->readMap<int>();
|
||||
auto box_ptr = boxes->readMap<float>();
|
||||
auto ids_ptr = ids->readMap<int>();
|
||||
auto score_ptr = scores->readMap<float>();
|
||||
|
||||
const int numResults = result_ids->getInfo()->size;
|
||||
for (int i = 0; i < numResults; i++) {
|
||||
int idx = result_ptr[i];
|
||||
if (idx < 0) break;
|
||||
float x1 = box_ptr[idx * 4 + 0] * scale;
|
||||
float y1 = box_ptr[idx * 4 + 1] * scale;
|
||||
float x2 = box_ptr[idx * 4 + 2] * scale;
|
||||
float y2 = box_ptr[idx * 4 + 3] * scale;
|
||||
int class_idx = ids_ptr[idx];
|
||||
float score = score_ptr[idx];
|
||||
printf("Detection: box = {%.2f, %.2f, %.2f, %.2f}, class = %s, score = %.2f\n",
|
||||
x1, y1, x2, y2, classNames[class_idx].c_str(), score);
|
||||
MNN::CV::rectangle(originalImage, { x1, y1 }, { x2, y2 }, { 0, 255, 0 }, 2);
|
||||
// Note: MNN::CV does not offer a putText function.
|
||||
// For text annotations, consider converting the image to cv::Mat and using OpenCV.
|
||||
}
|
||||
if (MNN::CV::imwrite("mnn_yolov8_cpp.jpg", originalImage)) {
|
||||
MNN_PRINT("Result image written to `mnn_yolov8_cpp.jpg`.\n");
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
MNN::Interpreter* interpreter;
|
||||
MNN::Session* session;
|
||||
MNN::Tensor* inputTensor;
|
||||
std::vector<int> inputDims;
|
||||
std::vector<std::string> classNames;
|
||||
};
|
||||
|
||||
int main(int argc, const char* argv[]) {
|
||||
if (argc < 3) {
|
||||
MNN_PRINT("Usage: ./main yolov8n.mnn input.jpg [backend] [precision] [thread]\n");
|
||||
return 0;
|
||||
}
|
||||
int backend = MNN_FORWARD_CPU;
|
||||
int precision = 1;
|
||||
int thread = 4;
|
||||
if (argc >= 4) {
|
||||
backend = atoi(argv[3]);
|
||||
}
|
||||
if (argc >= 5) {
|
||||
precision = atoi(argv[4]);
|
||||
}
|
||||
if (argc >= 6) {
|
||||
thread = atoi(argv[5]);
|
||||
}
|
||||
|
||||
Inference infer;
|
||||
if (!infer.loadModel(argv[1], backend, precision, thread))
|
||||
return 1;
|
||||
|
||||
const clock_t begin_time = clock();
|
||||
float scale = 1.0f;
|
||||
VARP originalImage = imread(argv[2]);
|
||||
VARP input = infer.preprocess(originalImage, 640, scale);
|
||||
auto preprocess_time = 1000.0 * (clock() - begin_time) / CLOCKS_PER_SEC;
|
||||
|
||||
const clock_t begin_time2 = clock();
|
||||
infer.runInference(input);
|
||||
auto inference_time = 1000.0 * (clock() - begin_time2) / CLOCKS_PER_SEC;
|
||||
const clock_t begin_time3 = clock();
|
||||
infer.postprocess(scale, originalImage);
|
||||
|
||||
auto postprocess_time = 1000.0 * (clock() - begin_time3) / CLOCKS_PER_SEC;
|
||||
printf("Speed: %.1fms preprocess, %.1fms inference, %.1fms postprocess\n",
|
||||
preprocess_time, inference_time, postprocess_time);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
cmake_minimum_required(VERSION 3.5)
|
||||
|
||||
set(PROJECT_NAME Yolov8OnnxRuntimeCPPInference)
|
||||
project(${PROJECT_NAME} VERSION 0.0.1 LANGUAGES CXX)
|
||||
|
||||
|
||||
# -------------- Support C++17 for using filesystem ------------------#
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
set(CMAKE_CXX_EXTENSIONS ON)
|
||||
set(CMAKE_INCLUDE_CURRENT_DIR ON)
|
||||
|
||||
|
||||
# -------------- OpenCV ------------------#
|
||||
find_package(OpenCV REQUIRED)
|
||||
include_directories(${OpenCV_INCLUDE_DIRS})
|
||||
|
||||
|
||||
# -------------- Compile CUDA for FP16 inference if needed ------------------#
|
||||
option(USE_CUDA "Enable CUDA support" ON)
|
||||
if (NOT APPLE AND USE_CUDA)
|
||||
find_package(CUDA REQUIRED)
|
||||
include_directories(${CUDA_INCLUDE_DIRS})
|
||||
add_definitions(-DUSE_CUDA)
|
||||
else ()
|
||||
set(USE_CUDA OFF)
|
||||
endif ()
|
||||
|
||||
# -------------- ONNXRUNTIME ------------------#
|
||||
|
||||
# Set ONNXRUNTIME_VERSION
|
||||
set(ONNXRUNTIME_VERSION 1.15.1)
|
||||
|
||||
if (NOT DEFINED ONNXRUNTIME_ROOT OR ONNXRUNTIME_ROOT STREQUAL "")
|
||||
if (WIN32)
|
||||
if (USE_CUDA)
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-win-x64-gpu-${ONNXRUNTIME_VERSION}")
|
||||
else ()
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-win-x64-${ONNXRUNTIME_VERSION}")
|
||||
endif ()
|
||||
elseif (LINUX)
|
||||
if (USE_CUDA)
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-linux-x64-gpu-${ONNXRUNTIME_VERSION}")
|
||||
else ()
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}")
|
||||
endif ()
|
||||
elseif (APPLE)
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-osx-arm64-${ONNXRUNTIME_VERSION}")
|
||||
# Apple X64 binary
|
||||
# set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-osx-x64-${ONNXRUNTIME_VERSION}")
|
||||
# Apple Universal binary
|
||||
# set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-osx-universal2-${ONNXRUNTIME_VERSION}")
|
||||
else ()
|
||||
message(SEND_ERROR "Variable ONNXRUNTIME_ROOT is not set properly. Please check if your cmake project \
|
||||
is not compiled with `-D WIN32=TRUE`, `-D LINUX=TRUE`, or `-D APPLE=TRUE`!")
|
||||
endif ()
|
||||
endif()
|
||||
|
||||
include_directories(${PROJECT_NAME} ${ONNXRUNTIME_ROOT}/include)
|
||||
|
||||
set(PROJECT_SOURCES
|
||||
main.cpp
|
||||
inference.h
|
||||
inference.cpp
|
||||
)
|
||||
|
||||
add_executable(${PROJECT_NAME} ${PROJECT_SOURCES})
|
||||
|
||||
if (WIN32)
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/onnxruntime.lib)
|
||||
if (USE_CUDA)
|
||||
target_link_libraries(${PROJECT_NAME} ${CUDA_LIBRARIES})
|
||||
endif ()
|
||||
elseif (LINUX)
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.so)
|
||||
if (USE_CUDA)
|
||||
target_link_libraries(${PROJECT_NAME} ${CUDA_LIBRARIES})
|
||||
endif ()
|
||||
elseif (APPLE)
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.dylib)
|
||||
endif ()
|
||||
|
||||
# For windows system, copy onnxruntime.dll to the same folder of the executable file
|
||||
if (WIN32)
|
||||
add_custom_command(TARGET ${PROJECT_NAME} POST_BUILD
|
||||
COMMAND ${CMAKE_COMMAND} -E copy_if_different
|
||||
"${ONNXRUNTIME_ROOT}/lib/onnxruntime.dll"
|
||||
$<TARGET_FILE_DIR:${PROJECT_NAME}>)
|
||||
endif ()
|
||||
|
||||
# Download https://raw.githubusercontent.com/ultralytics/ultralytics/main/ultralytics/cfg/datasets/coco.yaml
|
||||
# and put it in the same folder of the executable file
|
||||
configure_file(coco.yaml ${CMAKE_CURRENT_BINARY_DIR}/coco.yaml COPYONLY)
|
||||
|
||||
# Copy yolov8n.onnx file to the same folder of the executable file
|
||||
configure_file(yolov8n.onnx ${CMAKE_CURRENT_BINARY_DIR}/yolov8n.onnx COPYONLY)
|
||||
|
||||
# Create folder name images in the same folder of the executable file
|
||||
add_custom_command(TARGET ${PROJECT_NAME} POST_BUILD
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}/images
|
||||
)
|
||||
@@ -0,0 +1,173 @@
|
||||
# YOLOv8 ONNX Runtime C++ Example
|
||||
|
||||
<img alt="C++" src="https://img.shields.io/badge/C++-17-blue.svg?style=flat&logo=c%2B%2B"> <img alt="Onnx-runtime" src="https://img.shields.io/badge/OnnxRuntime-717272.svg?logo=Onnx&logoColor=white">
|
||||
|
||||
This example provides a practical guide on performing inference with [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) models using [C++](https://isocpp.org/), leveraging the capabilities of the [ONNX Runtime](https://onnxruntime.ai/) and the [OpenCV](https://opencv.org/) library. It's designed for developers looking to integrate YOLOv8 into C++ applications for efficient object detection.
|
||||
|
||||
## ✨ Benefits
|
||||
|
||||
- **Deployment-Friendly:** Well-suited for deployment in industrial and production environments.
|
||||
- **Performance:** Offers faster [inference latency](https://www.ultralytics.com/glossary/inference-latency) compared to OpenCV's DNN module on both CPU and [GPU](https://www.ultralytics.com/glossary/gpu-graphics-processing-unit).
|
||||
- **Acceleration:** Supports FP32 and [FP16 (Half Precision)](https://www.ultralytics.com/glossary/half-precision) inference acceleration using [NVIDIA CUDA](https://developer.nvidia.com/cuda-toolkit).
|
||||
|
||||
## ☕ Note
|
||||
|
||||
Thanks to recent updates in Ultralytics, YOLOv8 models now include a `Transpose` operation, aligning their output shape with YOLOv5. This allows the C++ code in this project to run inference seamlessly for YOLOv5, YOLOv7, and YOLOv8 models exported to the [ONNX format](https://onnx.ai/).
|
||||
|
||||
## 📦 Exporting YOLOv8 Models
|
||||
|
||||
You can export your trained [Ultralytics YOLO](https://docs.ultralytics.com/) models to the ONNX format required by this project. Use the Ultralytics `export` mode for this.
|
||||
|
||||
### Python
|
||||
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
|
||||
# Load a YOLOv8 model (e.g., yolov8n.pt)
|
||||
model = YOLO("yolov8n.pt")
|
||||
|
||||
# Export the model to ONNX format
|
||||
# opset=12 is recommended for compatibility
|
||||
# simplify=True optimizes the model graph
|
||||
# dynamic=False ensures fixed input size, often better for C++ deployment
|
||||
# imgsz=640 sets the input image size
|
||||
model.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=640)
|
||||
print("Model exported successfully to yolov8n.onnx")
|
||||
```
|
||||
|
||||
### CLI
|
||||
|
||||
```bash
|
||||
# Export the model using the command line
|
||||
yolo export model=yolov8n.pt format=onnx opset=12 simplify=True dynamic=False imgsz=640
|
||||
```
|
||||
|
||||
For more details on exporting models, refer to the [Ultralytics Export documentation](https://docs.ultralytics.com/modes/export/).
|
||||
|
||||
## 📦 Exporting YOLOv8 FP16 Models
|
||||
|
||||
To potentially gain further performance on compatible hardware (like NVIDIA GPUs), you can convert the exported FP32 ONNX model to FP16.
|
||||
|
||||
```python
|
||||
import onnx
|
||||
from onnxconverter_common import (
|
||||
float16,
|
||||
) # Ensure you have onnxconverter-common installed: pip install onnxconverter-common
|
||||
|
||||
# Load your FP32 ONNX model
|
||||
fp32_model_path = "yolov8n.onnx"
|
||||
model = onnx.load(fp32_model_path)
|
||||
|
||||
# Convert the model to FP16
|
||||
model_fp16 = float16.convert_float_to_float16(model)
|
||||
|
||||
# Save the FP16 model
|
||||
fp16_model_path = "yolov8n_fp16.onnx"
|
||||
onnx.save(model_fp16, fp16_model_path)
|
||||
print(f"Model converted and saved to {fp16_model_path}")
|
||||
```
|
||||
|
||||
## 📂 Download COCO YAML File
|
||||
|
||||
This example uses class names defined in a YAML file. You'll need the `coco.yaml` file, which corresponds to the standard [COCO dataset](https://docs.ultralytics.com/datasets/detect/coco/) classes. Download it directly:
|
||||
|
||||
- [Download coco.yaml](https://raw.githubusercontent.com/ultralytics/ultralytics/main/ultralytics/cfg/datasets/coco.yaml)
|
||||
|
||||
Save this file in the same directory where you plan to run the executable, or adjust the path in the C++ code accordingly.
|
||||
|
||||
## ⚙️ Dependencies
|
||||
|
||||
Ensure you have the following dependencies installed:
|
||||
|
||||
| Dependency | Version | Notes |
|
||||
| :------------------------------------------------------------------- | :------------ | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| [ONNX Runtime](https://onnxruntime.ai/docs/install/) | >=1.14.1 | Download pre-built binaries or build from source. Ensure GPU version if using CUDA. |
|
||||
| [OpenCV](https://opencv.org/releases/) | >=4.0.0 | Required for image loading and preprocessing. |
|
||||
| C++ Compiler | C++17 Support | Needed for features like `<filesystem>`. ([GCC](https://gcc.gnu.org/), [Clang](https://clang.llvm.org/), [MSVC](https://visualstudio.microsoft.com/vs/features/cplusplus/)) |
|
||||
| [CMake](https://cmake.org/download/) | >=3.18 | Cross-platform build system generator. Version 3.18+ recommended for better CUDA support discovery. |
|
||||
| [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit) (Optional) | >=11.4, <12.0 | Required for GPU acceleration via ONNX Runtime's CUDA Execution Provider. **Must be CUDA 11.x**. |
|
||||
| [cuDNN](https://developer.nvidia.com/cudnn) (CUDA required) | =8.x | Required by CUDA Execution Provider. **Must be cuDNN 8.x** compatible with your CUDA 11.x version. |
|
||||
|
||||
**Important Notes:**
|
||||
|
||||
1. **C++17:** The requirement stems from using the `<filesystem>` library introduced in C++17 for path handling.
|
||||
2. **CUDA/cuDNN Versions:** ONNX Runtime's CUDA execution provider currently has strict version requirements (CUDA 11.x, cuDNN 8.x). Check the latest [ONNX Runtime documentation](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html) for any updates to these constraints. Using incompatible versions will lead to runtime errors.
|
||||
|
||||
## 🛠️ Build Instructions
|
||||
|
||||
1. **Clone the Repository:**
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ultralytics/ultralytics.git
|
||||
cd ultralytics/examples/YOLOv8-ONNXRuntime-CPP
|
||||
```
|
||||
|
||||
2. **Create Build Directory:**
|
||||
|
||||
```bash
|
||||
mkdir build && cd build
|
||||
```
|
||||
|
||||
3. **Configure with CMake:**
|
||||
Run CMake to generate build files. You **must** specify the path to your ONNX Runtime installation directory using `ONNXRUNTIME_ROOT`. Adjust the path according to where you downloaded or built ONNX Runtime.
|
||||
|
||||
```bash
|
||||
# Example for Linux/macOS (adjust path as needed)
|
||||
cmake .. -DONNXRUNTIME_ROOT=/path/to/onnxruntime
|
||||
|
||||
# Example for Windows (adjust path as needed, use backslashes or forward slashes)
|
||||
cmake .. -DONNXRUNTIME_ROOT="C:/path/to/onnxruntime"
|
||||
```
|
||||
|
||||
**CMake Options:**
|
||||
- `-DONNXRUNTIME_ROOT=<path>`: **(Required)** Path to the extracted ONNX Runtime library.
|
||||
- `-DCMAKE_BUILD_TYPE=Release`: (Optional) Build in Release mode for optimizations.
|
||||
- If CMake struggles to find OpenCV, you might need to set `-DOpenCV_DIR=/path/to/opencv/build`.
|
||||
|
||||
4. **Build the Project:**
|
||||
Use the build tool generated by CMake (e.g., Make, Ninja, Visual Studio).
|
||||
|
||||
```bash
|
||||
# Using Make (common on Linux/macOS)
|
||||
make
|
||||
|
||||
# Using CMake's generic build command (works with Make, Ninja, etc.)
|
||||
cmake --build . --config Release
|
||||
```
|
||||
|
||||
5. **Locate Executable:**
|
||||
The compiled executable (e.g., `yolov8_onnxruntime_cpp`) will be located in the `build` directory.
|
||||
|
||||
## 🚀 Usage
|
||||
|
||||
Before running, ensure:
|
||||
|
||||
- The exported `.onnx` model file (e.g., `yolov8n.onnx`) is accessible.
|
||||
- The `coco.yaml` file is accessible.
|
||||
- Any required shared libraries for ONNX Runtime and OpenCV are in the system's PATH or accessible by the executable.
|
||||
|
||||
Modify the `main.cpp` file (or create a configuration mechanism) to set the parameters:
|
||||
|
||||
```cpp
|
||||
//change your param as you like
|
||||
//Pay attention to your device and the onnx model type(fp32 or fp16)
|
||||
DL_INIT_PARAM params;
|
||||
params.rectConfidenceThreshold = 0.1;
|
||||
params.iouThreshold = 0.5;
|
||||
params.modelPath = "yolov8n.onnx";
|
||||
params.imgSize = { 640, 640 };
|
||||
params.cudaEnable = true;
|
||||
params.modelType = YOLO_DETECT_V8;
|
||||
yoloDetector->CreateSession(params);
|
||||
Detector(yoloDetector);
|
||||
```
|
||||
|
||||
Run the executable from the `build` directory:
|
||||
|
||||
```bash
|
||||
./yolov8_onnxruntime_cpp
|
||||
```
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! If you find any issues or have suggestions for improvements, please feel free to open an issue or submit a pull request on the main [Ultralytics repository](https://github.com/ultralytics/ultralytics).
|
||||
@@ -0,0 +1,374 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include "inference.h"
|
||||
#include <regex>
|
||||
|
||||
#define benchmark
|
||||
#define min(a,b) (((a) < (b)) ? (a) : (b))
|
||||
YOLO_V8::YOLO_V8() {
|
||||
|
||||
}
|
||||
|
||||
|
||||
YOLO_V8::~YOLO_V8() {
|
||||
delete session;
|
||||
}
|
||||
|
||||
#ifdef USE_CUDA
|
||||
namespace Ort
|
||||
{
|
||||
template<>
|
||||
struct TypeToTensorType<half> { static constexpr ONNXTensorElementDataType type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; };
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
template<typename T>
|
||||
char* BlobFromImage(cv::Mat& iImg, T& iBlob) {
|
||||
int channels = iImg.channels();
|
||||
int imgHeight = iImg.rows;
|
||||
int imgWidth = iImg.cols;
|
||||
|
||||
for (int c = 0; c < channels; c++)
|
||||
{
|
||||
for (int h = 0; h < imgHeight; h++)
|
||||
{
|
||||
for (int w = 0; w < imgWidth; w++)
|
||||
{
|
||||
iBlob[c * imgWidth * imgHeight + h * imgWidth + w] = typename std::remove_pointer<T>::type(
|
||||
(iImg.at<cv::Vec3b>(h, w)[c]) / 255.0f);
|
||||
}
|
||||
}
|
||||
}
|
||||
return RET_OK;
|
||||
}
|
||||
|
||||
|
||||
char* YOLO_V8::PreProcess(cv::Mat& iImg, std::vector<int> iImgSize, cv::Mat& oImg)
|
||||
{
|
||||
if (iImg.channels() == 3)
|
||||
{
|
||||
oImg = iImg.clone();
|
||||
cv::cvtColor(oImg, oImg, cv::COLOR_BGR2RGB);
|
||||
}
|
||||
else
|
||||
{
|
||||
cv::cvtColor(iImg, oImg, cv::COLOR_GRAY2RGB);
|
||||
}
|
||||
|
||||
switch (modelType)
|
||||
{
|
||||
case YOLO_DETECT_V8:
|
||||
case YOLO_POSE:
|
||||
case YOLO_DETECT_V8_HALF:
|
||||
case YOLO_POSE_V8_HALF://LetterBox
|
||||
{
|
||||
int new_h = iImgSize.at(0);
|
||||
int new_w = iImgSize.at(1);
|
||||
float r = min(new_w / (float)iImg.cols, new_h / (float)iImg.rows);
|
||||
int resized_w = static_cast<int>(iImg.cols * r);
|
||||
int resized_h = static_cast<int>(iImg.rows * r);
|
||||
resizeScales = 1.0f / r;
|
||||
cv::resize(oImg, oImg, cv::Size(resized_w, resized_h));
|
||||
cv::Mat tempImg = cv::Mat::zeros(new_h, new_w, CV_8UC3);
|
||||
oImg.copyTo(tempImg(cv::Rect(0, 0, resized_w, resized_h)));
|
||||
oImg = tempImg;
|
||||
break;
|
||||
}
|
||||
case YOLO_CLS://CenterCrop
|
||||
{
|
||||
int h = iImg.rows;
|
||||
int w = iImg.cols;
|
||||
int m = min(h, w);
|
||||
int top = (h - m) / 2;
|
||||
int left = (w - m) / 2;
|
||||
cv::resize(oImg(cv::Rect(left, top, m, m)), oImg, cv::Size(iImgSize.at(1), iImgSize.at(0)));
|
||||
break;
|
||||
}
|
||||
}
|
||||
return RET_OK;
|
||||
}
|
||||
|
||||
|
||||
char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) {
|
||||
char* Ret = RET_OK;
|
||||
std::regex pattern("[\u4e00-\u9fa5]");
|
||||
bool result = std::regex_search(iParams.modelPath, pattern);
|
||||
if (result)
|
||||
{
|
||||
Ret = "[YOLO_V8]:Your model path is error.Change your model path without chinese characters.";
|
||||
std::cout << Ret << std::endl;
|
||||
return Ret;
|
||||
}
|
||||
try
|
||||
{
|
||||
rectConfidenceThreshold = iParams.rectConfidenceThreshold;
|
||||
iouThreshold = iParams.iouThreshold;
|
||||
imgSize = iParams.imgSize;
|
||||
modelType = iParams.modelType;
|
||||
cudaEnable = iParams.cudaEnable;
|
||||
env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "Yolo");
|
||||
Ort::SessionOptions sessionOption;
|
||||
if (iParams.cudaEnable)
|
||||
{
|
||||
OrtCUDAProviderOptions cudaOption;
|
||||
cudaOption.device_id = 0;
|
||||
sessionOption.AppendExecutionProvider_CUDA(cudaOption);
|
||||
}
|
||||
sessionOption.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
sessionOption.SetIntraOpNumThreads(iParams.intraOpNumThreads);
|
||||
sessionOption.SetLogSeverityLevel(iParams.logSeverityLevel);
|
||||
|
||||
#ifdef _WIN32
|
||||
int ModelPathSize = MultiByteToWideChar(CP_UTF8, 0, iParams.modelPath.c_str(), static_cast<int>(iParams.modelPath.length()), nullptr, 0);
|
||||
wchar_t* wide_cstr = new wchar_t[ModelPathSize + 1];
|
||||
MultiByteToWideChar(CP_UTF8, 0, iParams.modelPath.c_str(), static_cast<int>(iParams.modelPath.length()), wide_cstr, ModelPathSize);
|
||||
wide_cstr[ModelPathSize] = L'\0';
|
||||
const wchar_t* modelPath = wide_cstr;
|
||||
#else
|
||||
const char* modelPath = iParams.modelPath.c_str();
|
||||
#endif // _WIN32
|
||||
|
||||
session = new Ort::Session(env, modelPath, sessionOption);
|
||||
Ort::AllocatorWithDefaultOptions allocator;
|
||||
size_t inputNodesNum = session->GetInputCount();
|
||||
for (size_t i = 0; i < inputNodesNum; i++)
|
||||
{
|
||||
Ort::AllocatedStringPtr input_node_name = session->GetInputNameAllocated(i, allocator);
|
||||
char* temp_buf = new char[50];
|
||||
strcpy(temp_buf, input_node_name.get());
|
||||
inputNodeNames.push_back(temp_buf);
|
||||
}
|
||||
size_t OutputNodesNum = session->GetOutputCount();
|
||||
for (size_t i = 0; i < OutputNodesNum; i++)
|
||||
{
|
||||
Ort::AllocatedStringPtr output_node_name = session->GetOutputNameAllocated(i, allocator);
|
||||
char* temp_buf = new char[10];
|
||||
strcpy(temp_buf, output_node_name.get());
|
||||
outputNodeNames.push_back(temp_buf);
|
||||
}
|
||||
options = Ort::RunOptions{ nullptr };
|
||||
WarmUpSession();
|
||||
return RET_OK;
|
||||
}
|
||||
catch (const std::exception& e)
|
||||
{
|
||||
const char* str1 = "[YOLO_V8]:";
|
||||
const char* str2 = e.what();
|
||||
std::string result = std::string(str1) + std::string(str2);
|
||||
char* merged = new char[result.length() + 1];
|
||||
std::strcpy(merged, result.c_str());
|
||||
std::cout << merged << std::endl;
|
||||
delete[] merged;
|
||||
return "[YOLO_V8]:Create session failed.";
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
char* YOLO_V8::RunSession(cv::Mat& iImg, std::vector<DL_RESULT>& oResult) {
|
||||
#ifdef benchmark
|
||||
clock_t starttime_1 = clock();
|
||||
#endif // benchmark
|
||||
|
||||
char* Ret = RET_OK;
|
||||
cv::Mat processedImg;
|
||||
PreProcess(iImg, imgSize, processedImg);
|
||||
if (modelType < 4)
|
||||
{
|
||||
float* blob = new float[processedImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> inputNodeDims = { 1, 3, imgSize.at(0), imgSize.at(1) };
|
||||
TensorProcess(starttime_1, iImg, blob, inputNodeDims, oResult);
|
||||
}
|
||||
else
|
||||
{
|
||||
#ifdef USE_CUDA
|
||||
half* blob = new half[processedImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> inputNodeDims = { 1,3,imgSize.at(0),imgSize.at(1) };
|
||||
TensorProcess(starttime_1, iImg, blob, inputNodeDims, oResult);
|
||||
#endif
|
||||
}
|
||||
|
||||
return Ret;
|
||||
}
|
||||
|
||||
|
||||
template<typename N>
|
||||
char* YOLO_V8::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std::vector<int64_t>& inputNodeDims,
|
||||
std::vector<DL_RESULT>& oResult) {
|
||||
Ort::Value inputTensor = Ort::Value::CreateTensor<typename std::remove_pointer<N>::type>(
|
||||
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1),
|
||||
inputNodeDims.data(), inputNodeDims.size());
|
||||
#ifdef benchmark
|
||||
clock_t starttime_2 = clock();
|
||||
#endif // benchmark
|
||||
auto outputTensor = session->Run(options, inputNodeNames.data(), &inputTensor, 1, outputNodeNames.data(),
|
||||
outputNodeNames.size());
|
||||
#ifdef benchmark
|
||||
clock_t starttime_3 = clock();
|
||||
#endif // benchmark
|
||||
|
||||
Ort::TypeInfo typeInfo = outputTensor.front().GetTypeInfo();
|
||||
auto tensor_info = typeInfo.GetTensorTypeAndShapeInfo();
|
||||
std::vector<int64_t> outputNodeDims = tensor_info.GetShape();
|
||||
auto output = outputTensor.front().GetTensorMutableData<typename std::remove_pointer<N>::type>();
|
||||
delete[] blob;
|
||||
switch (modelType)
|
||||
{
|
||||
case YOLO_DETECT_V8:
|
||||
case YOLO_DETECT_V8_HALF:
|
||||
{
|
||||
int signalResultNum = outputNodeDims[1];//84
|
||||
int strideNum = outputNodeDims[2];//8400
|
||||
std::vector<int> class_ids;
|
||||
std::vector<float> confidences;
|
||||
std::vector<cv::Rect> boxes;
|
||||
cv::Mat rawData;
|
||||
if (modelType == YOLO_DETECT_V8)
|
||||
{
|
||||
// FP32
|
||||
rawData = cv::Mat(signalResultNum, strideNum, CV_32F, output);
|
||||
}
|
||||
else
|
||||
{
|
||||
// FP16
|
||||
rawData = cv::Mat(signalResultNum, strideNum, CV_16F, output);
|
||||
rawData.convertTo(rawData, CV_32F);
|
||||
}
|
||||
// Note:
|
||||
// ultralytics add transpose operator to the output of yolov8 model.which make yolov8/v5/v7 has same shape
|
||||
// https://github.com/ultralytics/assets/releases/download/v8.4.0/yolov8n.pt
|
||||
rawData = rawData.t();
|
||||
|
||||
float* data = (float*)rawData.data;
|
||||
|
||||
for (int i = 0; i < strideNum; ++i)
|
||||
{
|
||||
float* classesScores = data + 4;
|
||||
cv::Mat scores(1, this->classes.size(), CV_32FC1, classesScores);
|
||||
cv::Point class_id;
|
||||
double maxClassScore;
|
||||
cv::minMaxLoc(scores, 0, &maxClassScore, 0, &class_id);
|
||||
if (maxClassScore > rectConfidenceThreshold)
|
||||
{
|
||||
confidences.push_back(maxClassScore);
|
||||
class_ids.push_back(class_id.x);
|
||||
float x = data[0];
|
||||
float y = data[1];
|
||||
float w = data[2];
|
||||
float h = data[3];
|
||||
|
||||
int left = int((x - 0.5 * w) * resizeScales);
|
||||
int top = int((y - 0.5 * h) * resizeScales);
|
||||
|
||||
int width = int(w * resizeScales);
|
||||
int height = int(h * resizeScales);
|
||||
|
||||
boxes.push_back(cv::Rect(left, top, width, height));
|
||||
}
|
||||
data += signalResultNum;
|
||||
}
|
||||
std::vector<int> nmsResult;
|
||||
cv::dnn::NMSBoxes(boxes, confidences, rectConfidenceThreshold, iouThreshold, nmsResult);
|
||||
for (int i = 0; i < nmsResult.size(); ++i)
|
||||
{
|
||||
int idx = nmsResult[i];
|
||||
DL_RESULT result;
|
||||
result.classId = class_ids[idx];
|
||||
result.confidence = confidences[idx];
|
||||
result.box = boxes[idx];
|
||||
oResult.push_back(result);
|
||||
}
|
||||
|
||||
#ifdef benchmark
|
||||
clock_t starttime_4 = clock();
|
||||
double pre_process_time = (double)(starttime_2 - starttime_1) / CLOCKS_PER_SEC * 1000;
|
||||
double process_time = (double)(starttime_3 - starttime_2) / CLOCKS_PER_SEC * 1000;
|
||||
double post_process_time = (double)(starttime_4 - starttime_3) / CLOCKS_PER_SEC * 1000;
|
||||
if (cudaEnable)
|
||||
{
|
||||
std::cout << "[YOLO_V8(CUDA)]: " << pre_process_time << "ms pre-process, " << process_time << "ms inference, " << post_process_time << "ms post-process." << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "[YOLO_V8(CPU)]: " << pre_process_time << "ms pre-process, " << process_time << "ms inference, " << post_process_time << "ms post-process." << std::endl;
|
||||
}
|
||||
#endif // benchmark
|
||||
|
||||
break;
|
||||
}
|
||||
case YOLO_CLS:
|
||||
case YOLO_CLS_HALF:
|
||||
{
|
||||
cv::Mat rawData;
|
||||
if (modelType == YOLO_CLS) {
|
||||
// FP32
|
||||
rawData = cv::Mat(1, this->classes.size(), CV_32F, output);
|
||||
} else {
|
||||
// FP16
|
||||
rawData = cv::Mat(1, this->classes.size(), CV_16F, output);
|
||||
rawData.convertTo(rawData, CV_32F);
|
||||
}
|
||||
float *data = (float *) rawData.data;
|
||||
|
||||
DL_RESULT result;
|
||||
for (int i = 0; i < this->classes.size(); i++)
|
||||
{
|
||||
result.classId = i;
|
||||
result.confidence = data[i];
|
||||
oResult.push_back(result);
|
||||
}
|
||||
break;
|
||||
}
|
||||
default:
|
||||
std::cout << "[YOLO_V8]: " << "Not support model type." << std::endl;
|
||||
}
|
||||
return RET_OK;
|
||||
|
||||
}
|
||||
|
||||
|
||||
char* YOLO_V8::WarmUpSession() {
|
||||
clock_t starttime_1 = clock();
|
||||
cv::Mat iImg = cv::Mat(cv::Size(imgSize.at(1), imgSize.at(0)), CV_8UC3);
|
||||
cv::Mat processedImg;
|
||||
PreProcess(iImg, imgSize, processedImg);
|
||||
if (modelType < 4)
|
||||
{
|
||||
float* blob = new float[iImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> YOLO_input_node_dims = { 1, 3, imgSize.at(0), imgSize.at(1) };
|
||||
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
|
||||
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1),
|
||||
YOLO_input_node_dims.data(), YOLO_input_node_dims.size());
|
||||
auto output_tensors = session->Run(options, inputNodeNames.data(), &input_tensor, 1, outputNodeNames.data(),
|
||||
outputNodeNames.size());
|
||||
delete[] blob;
|
||||
clock_t starttime_4 = clock();
|
||||
double post_process_time = (double)(starttime_4 - starttime_1) / CLOCKS_PER_SEC * 1000;
|
||||
if (cudaEnable)
|
||||
{
|
||||
std::cout << "[YOLO_V8(CUDA)]: " << "Cuda warm-up cost " << post_process_time << " ms. " << std::endl;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
#ifdef USE_CUDA
|
||||
half* blob = new half[iImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> YOLO_input_node_dims = { 1,3,imgSize.at(0),imgSize.at(1) };
|
||||
Ort::Value input_tensor = Ort::Value::CreateTensor<half>(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1), YOLO_input_node_dims.data(), YOLO_input_node_dims.size());
|
||||
auto output_tensors = session->Run(options, inputNodeNames.data(), &input_tensor, 1, outputNodeNames.data(), outputNodeNames.size());
|
||||
delete[] blob;
|
||||
clock_t starttime_4 = clock();
|
||||
double post_process_time = (double)(starttime_4 - starttime_1) / CLOCKS_PER_SEC * 1000;
|
||||
if (cudaEnable)
|
||||
{
|
||||
std::cout << "[YOLO_V8(CUDA)]: " << "Cuda warm-up cost " << post_process_time << " ms. " << std::endl;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
return RET_OK;
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#pragma once
|
||||
|
||||
#define RET_OK nullptr
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <Windows.h>
|
||||
#include <direct.h>
|
||||
#include <io.h>
|
||||
#endif
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <cstdio>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include "onnxruntime_cxx_api.h"
|
||||
|
||||
#ifdef USE_CUDA
|
||||
#include <cuda_fp16.h>
|
||||
#endif
|
||||
|
||||
|
||||
enum MODEL_TYPE
|
||||
{
|
||||
//FLOAT32 MODEL
|
||||
YOLO_DETECT_V8 = 1,
|
||||
YOLO_POSE = 2,
|
||||
YOLO_CLS = 3,
|
||||
|
||||
//FLOAT16 MODEL
|
||||
YOLO_DETECT_V8_HALF = 4,
|
||||
YOLO_POSE_V8_HALF = 5,
|
||||
YOLO_CLS_HALF = 6
|
||||
};
|
||||
|
||||
|
||||
typedef struct _DL_INIT_PARAM
|
||||
{
|
||||
std::string modelPath;
|
||||
MODEL_TYPE modelType = YOLO_DETECT_V8;
|
||||
std::vector<int> imgSize = { 640, 640 };
|
||||
float rectConfidenceThreshold = 0.6;
|
||||
float iouThreshold = 0.5;
|
||||
int keyPointsNum = 2;//Note:kpt number for pose
|
||||
bool cudaEnable = false;
|
||||
int logSeverityLevel = 3;
|
||||
int intraOpNumThreads = 1;
|
||||
} DL_INIT_PARAM;
|
||||
|
||||
|
||||
typedef struct _DL_RESULT
|
||||
{
|
||||
int classId;
|
||||
float confidence;
|
||||
cv::Rect box;
|
||||
std::vector<cv::Point2f> keyPoints;
|
||||
} DL_RESULT;
|
||||
|
||||
|
||||
class YOLO_V8
|
||||
{
|
||||
public:
|
||||
YOLO_V8();
|
||||
|
||||
~YOLO_V8();
|
||||
|
||||
public:
|
||||
char* CreateSession(DL_INIT_PARAM& iParams);
|
||||
|
||||
char* RunSession(cv::Mat& iImg, std::vector<DL_RESULT>& oResult);
|
||||
|
||||
char* WarmUpSession();
|
||||
|
||||
template<typename N>
|
||||
char* TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std::vector<int64_t>& inputNodeDims,
|
||||
std::vector<DL_RESULT>& oResult);
|
||||
|
||||
char* PreProcess(cv::Mat& iImg, std::vector<int> iImgSize, cv::Mat& oImg);
|
||||
|
||||
std::vector<std::string> classes{};
|
||||
|
||||
private:
|
||||
Ort::Env env;
|
||||
Ort::Session* session;
|
||||
bool cudaEnable;
|
||||
Ort::RunOptions options;
|
||||
std::vector<const char*> inputNodeNames;
|
||||
std::vector<const char*> outputNodeNames;
|
||||
|
||||
MODEL_TYPE modelType;
|
||||
std::vector<int> imgSize;
|
||||
float rectConfidenceThreshold;
|
||||
float iouThreshold;
|
||||
float resizeScales;//letterbox scale
|
||||
};
|
||||
@@ -0,0 +1,195 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#include <iostream>
|
||||
#include <iomanip>
|
||||
#include "inference.h"
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <random>
|
||||
|
||||
void Detector(YOLO_V8*& p) {
|
||||
std::filesystem::path current_path = std::filesystem::current_path();
|
||||
std::filesystem::path imgs_path = current_path / "images";
|
||||
for (auto& i : std::filesystem::directory_iterator(imgs_path))
|
||||
{
|
||||
if (i.path().extension() == ".jpg" || i.path().extension() == ".png" || i.path().extension() == ".jpeg")
|
||||
{
|
||||
std::string img_path = i.path().string();
|
||||
cv::Mat img = cv::imread(img_path);
|
||||
std::vector<DL_RESULT> res;
|
||||
p->RunSession(img, res);
|
||||
|
||||
for (auto& re : res)
|
||||
{
|
||||
cv::RNG rng(cv::getTickCount());
|
||||
cv::Scalar color(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256));
|
||||
|
||||
cv::rectangle(img, re.box, color, 3);
|
||||
|
||||
float confidence = floor(100 * re.confidence) / 100;
|
||||
std::cout << std::fixed << std::setprecision(2);
|
||||
std::string label = p->classes[re.classId] + " " +
|
||||
std::to_string(confidence).substr(0, std::to_string(confidence).size() - 4);
|
||||
|
||||
cv::rectangle(
|
||||
img,
|
||||
cv::Point(re.box.x, re.box.y - 25),
|
||||
cv::Point(re.box.x + label.length() * 15, re.box.y),
|
||||
color,
|
||||
cv::FILLED
|
||||
);
|
||||
|
||||
cv::putText(
|
||||
img,
|
||||
label,
|
||||
cv::Point(re.box.x, re.box.y - 5),
|
||||
cv::FONT_HERSHEY_SIMPLEX,
|
||||
0.75,
|
||||
cv::Scalar(0, 0, 0),
|
||||
2
|
||||
);
|
||||
|
||||
|
||||
}
|
||||
std::cout << "Press any key to exit" << std::endl;
|
||||
cv::imshow("Result of Detection", img);
|
||||
cv::waitKey(0);
|
||||
cv::destroyAllWindows();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void Classifier(YOLO_V8*& p)
|
||||
{
|
||||
std::filesystem::path current_path = std::filesystem::current_path();
|
||||
std::filesystem::path imgs_path = current_path;// / "images"
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_int_distribution<int> dis(0, 255);
|
||||
for (auto& i : std::filesystem::directory_iterator(imgs_path))
|
||||
{
|
||||
if (i.path().extension() == ".jpg" || i.path().extension() == ".png")
|
||||
{
|
||||
std::string img_path = i.path().string();
|
||||
//std::cout << img_path << std::endl;
|
||||
cv::Mat img = cv::imread(img_path);
|
||||
std::vector<DL_RESULT> res;
|
||||
char* ret = p->RunSession(img, res);
|
||||
|
||||
float positionY = 50;
|
||||
for (int i = 0; i < res.size(); i++)
|
||||
{
|
||||
int r = dis(gen);
|
||||
int g = dis(gen);
|
||||
int b = dis(gen);
|
||||
cv::putText(img, std::to_string(i) + ":", cv::Point(10, positionY), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2);
|
||||
cv::putText(img, std::to_string(res.at(i).confidence), cv::Point(70, positionY), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2);
|
||||
positionY += 50;
|
||||
}
|
||||
|
||||
cv::imshow("TEST_CLS", img);
|
||||
cv::waitKey(0);
|
||||
cv::destroyAllWindows();
|
||||
//cv::imwrite("E:\\output\\" + std::to_string(k) + ".png", img);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
int ReadCocoYaml(YOLO_V8*& p) {
|
||||
// Open the YAML file
|
||||
std::ifstream file("coco.yaml");
|
||||
if (!file.is_open())
|
||||
{
|
||||
std::cerr << "Failed to open file" << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Read the file line by line
|
||||
std::string line;
|
||||
std::vector<std::string> lines;
|
||||
while (std::getline(file, line))
|
||||
{
|
||||
lines.push_back(line);
|
||||
}
|
||||
|
||||
// Find the start and end of the names section
|
||||
std::size_t start = 0;
|
||||
std::size_t end = 0;
|
||||
for (std::size_t i = 0; i < lines.size(); i++)
|
||||
{
|
||||
if (lines[i].find("names:") != std::string::npos)
|
||||
{
|
||||
start = i + 1;
|
||||
}
|
||||
else if (start > 0 && lines[i].find(':') == std::string::npos)
|
||||
{
|
||||
end = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Extract the names
|
||||
std::vector<std::string> names;
|
||||
for (std::size_t i = start; i < end; i++)
|
||||
{
|
||||
std::stringstream ss(lines[i]);
|
||||
std::string name;
|
||||
std::getline(ss, name, ':'); // Extract the number before the delimiter
|
||||
std::getline(ss, name); // Extract the string after the delimiter
|
||||
names.push_back(name);
|
||||
}
|
||||
|
||||
p->classes = names;
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
void DetectTest()
|
||||
{
|
||||
YOLO_V8* yoloDetector = new YOLO_V8;
|
||||
ReadCocoYaml(yoloDetector);
|
||||
DL_INIT_PARAM params;
|
||||
params.rectConfidenceThreshold = 0.1;
|
||||
params.iouThreshold = 0.5;
|
||||
params.modelPath = "yolov8n.onnx";
|
||||
params.imgSize = { 640, 640 };
|
||||
#ifdef USE_CUDA
|
||||
params.cudaEnable = true;
|
||||
|
||||
// GPU FP32 inference
|
||||
params.modelType = YOLO_DETECT_V8;
|
||||
// GPU FP16 inference
|
||||
//Note: change fp16 onnx model
|
||||
//params.modelType = YOLO_DETECT_V8_HALF;
|
||||
|
||||
#else
|
||||
// CPU inference
|
||||
params.modelType = YOLO_DETECT_V8;
|
||||
params.cudaEnable = false;
|
||||
|
||||
#endif
|
||||
yoloDetector->CreateSession(params);
|
||||
Detector(yoloDetector);
|
||||
}
|
||||
|
||||
|
||||
void ClsTest()
|
||||
{
|
||||
YOLO_V8* yoloDetector = new YOLO_V8;
|
||||
std::string model_path = "cls.onnx";
|
||||
ReadCocoYaml(yoloDetector);
|
||||
DL_INIT_PARAM params{ model_path, YOLO_CLS, {224, 224} };
|
||||
yoloDetector->CreateSession(params);
|
||||
Classifier(yoloDetector);
|
||||
}
|
||||
|
||||
|
||||
int main()
|
||||
{
|
||||
//DetectTest();
|
||||
ClsTest();
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
[package]
|
||||
name = "yolov8-rs"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
|
||||
[dependencies]
|
||||
clap = { version = "4.2.4", features = ["derive"] }
|
||||
image = { version = "0.25.2"}
|
||||
imageproc = { version = "0.25.0"}
|
||||
ndarray = { version = "0.16" }
|
||||
ort = { version = "2.0.0-rc.5", features = ["cuda", "tensorrt", "download-binaries", "copy-dylibs", "half"]}
|
||||
rusttype = { version = "0.9.3" }
|
||||
anyhow = { version = "1.0.75" }
|
||||
regex = { version = "1.5.4" }
|
||||
rand = { version = "0.8.5" }
|
||||
chrono = { version = "0.4.30" }
|
||||
half = { version = "2.3.1" }
|
||||
dirs = { version = "5.0.1" }
|
||||
ureq = { version = "2.9.1" }
|
||||
ab_glyph = "0.2.29"
|
||||
@@ -0,0 +1,230 @@
|
||||
# YOLOv8-ONNXRuntime-Rust for All Key YOLO Tasks
|
||||
|
||||
This repository provides a Rust demonstration for performing Ultralytics YOLOv8 tasks like [Classification](https://docs.ultralytics.com/tasks/classify/), [Segmentation](https://docs.ultralytics.com/tasks/segment/), [Detection](https://docs.ultralytics.com/tasks/detect/), [Pose Estimation](https://docs.ultralytics.com/tasks/pose/), and [Oriented Bounding Box (OBB)](https://docs.ultralytics.com/tasks/obb/) detection using the [ONNXRuntime](https://onnxruntime.ai/).
|
||||
|
||||
## ✨ Recently Updated
|
||||
|
||||
- Added YOLOv8-OBB demo.
|
||||
- Updated ONNXRuntime dependency to 1.19.x.
|
||||
|
||||
Newly updated YOLOv8 example code is located in [this repository](https://github.com/jamjamjon/usls/tree/main/examples/yolo).
|
||||
|
||||
## 🚀 Features
|
||||
|
||||
- Supports `Classification`, `Segmentation`, `Detection`, `Pose(Keypoints)-Detection`, and `OBB` tasks.
|
||||
- Supports `FP16` & `FP32` [ONNX](https://onnx.ai/) models.
|
||||
- Supports `CPU`, `CUDA`, and `TensorRT` execution providers to accelerate computation.
|
||||
- Supports dynamic input shapes (`batch`, `width`, `height`).
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
### 1. Install Rust
|
||||
|
||||
Please follow the official Rust installation guide: [https://www.rust-lang.org/tools/install](https://rust-lang.org/tools/install/).
|
||||
|
||||
### 2. ONNXRuntime Linking
|
||||
|
||||
- #### For detailed setup instructions, refer to the [ORT documentation](https://ort.pyke.io/setup/linking).
|
||||
|
||||
- #### For Linux or macOS Users:
|
||||
- Download the ONNX Runtime package from the [Releases page](https://github.com/microsoft/onnxruntime/releases).
|
||||
- Set up the library path by exporting the `ORT_DYLIB_PATH` environment variable:
|
||||
```bash
|
||||
export ORT_DYLIB_PATH=/path/to/onnxruntime/lib/libonnxruntime.so.1.19.0 # Adjust version/path as needed
|
||||
```
|
||||
|
||||
### 3. [Optional] Install CUDA & CuDNN & TensorRT
|
||||
|
||||
- The CUDA execution provider requires [CUDA](https://developer.nvidia.com/cuda-toolkit) v11.6+.
|
||||
- The TensorRT execution provider requires CUDA v11.4+ and [TensorRT](https://developer.nvidia.com/tensorrt) v8.4+. You may also need [cuDNN](https://developer.nvidia.com/cudnn).
|
||||
|
||||
## ▶️ Get Started
|
||||
|
||||
### 1. Export the Ultralytics YOLOv8 ONNX Models
|
||||
|
||||
First, install the Ultralytics package:
|
||||
|
||||
```bash
|
||||
pip install -U ultralytics
|
||||
```
|
||||
|
||||
Then, export the desired [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) models to the ONNX format. See the [Export documentation](https://docs.ultralytics.com/modes/export/) for more details.
|
||||
|
||||
```bash
|
||||
# Export ONNX model with dynamic shapes (recommended for flexibility)
|
||||
yolo export model=yolov8m.pt format=onnx simplify dynamic
|
||||
yolo export model=yolov8m-cls.pt format=onnx simplify dynamic
|
||||
yolo export model=yolov8m-pose.pt format=onnx simplify dynamic
|
||||
yolo export model=yolov8m-seg.pt format=onnx simplify dynamic
|
||||
# yolo export model=yolov8m-obb.pt format=onnx simplify dynamic # Add OBB export if needed
|
||||
|
||||
# Export ONNX model with constant shapes (if dynamic shapes are not required)
|
||||
# yolo export model=yolov8m.pt format=onnx simplify
|
||||
# yolo export model=yolov8m-cls.pt format=onnx simplify
|
||||
# yolo export model=yolov8m-pose.pt format=onnx simplify
|
||||
# yolo export model=yolov8m-seg.pt format=onnx simplify
|
||||
# yolo export model=yolov8m-obb.pt format=onnx simplify
|
||||
```
|
||||
|
||||
### 2. Run Inference
|
||||
|
||||
This command will perform inference using the specified ONNX model on the source image using the CPU.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --model MODEL_PATH.onnx --source SOURCE_IMAGE.jpg
|
||||
```
|
||||
|
||||
#### Using GPU Acceleration
|
||||
|
||||
Set `--cuda` to use the CUDA execution provider for faster inference on NVIDIA GPUs.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --cuda --model MODEL_PATH.onnx --source SOURCE_IMAGE.jpg
|
||||
```
|
||||
|
||||
Set `--trt` to use the TensorRT execution provider. You can also set `--fp16` simultaneously to leverage the TensorRT FP16 engine for potentially even greater speed, especially on compatible hardware.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --trt --fp16 --model MODEL_PATH.onnx --source SOURCE_IMAGE.jpg
|
||||
```
|
||||
|
||||
#### Specifying Device and Batch Size
|
||||
|
||||
Set `--device_id` to select a specific GPU device. If the specified device ID is invalid (e.g., setting `device_id 1` when only one GPU exists), `ort` will automatically fall back to the `CPU` execution provider without causing a panic.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --cuda --device_id 0 --model MODEL_PATH.onnx --source SOURCE_IMAGE.jpg
|
||||
```
|
||||
|
||||
Set `--batch` to perform inference with a specific batch size.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --cuda --batch 2 --model MODEL_PATH.onnx --source SOURCE_IMAGE.jpg
|
||||
```
|
||||
|
||||
If you're using `--trt` with a model exported with dynamic batch dimensions, you can explicitly specify the minimum, optimal, and maximum batch sizes for TensorRT optimization using `--batch-min`, `--batch`, and `--batch-max`. Refer to the [TensorRT Execution Provider documentation](https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html#explicit-shape-range-for-dynamic-shape-input) for details.
|
||||
|
||||
#### Dynamic Image Size
|
||||
|
||||
Set `--height` and `--width` to perform inference with dynamic image sizes. **Note:** The ONNX model must have been exported with dynamic input shapes (`dynamic=True`).
|
||||
|
||||
```bash
|
||||
cargo run --release -- --cuda --width 480 --height 640 --model MODEL_PATH_dynamic.onnx --source SOURCE_IMAGE.jpg
|
||||
```
|
||||
|
||||
#### Profiling Performance
|
||||
|
||||
Set `--profile` to measure the time consumed in each stage of the inference pipeline (preprocessing, H2D transfer, inference, D2H transfer, postprocessing). **Note:** Models often require a few "warm-up" runs (1-3 iterations) before reaching optimal performance. Ensure you run the command enough times to get a stable performance evaluation.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --trt --fp16 --profile --model MODEL_PATH.onnx --source SOURCE_IMAGE.jpg
|
||||
```
|
||||
|
||||
Example Profile Output (yolov8m.onnx, batch=1, 3 runs, trt, fp16, RTX 3060Ti):
|
||||
|
||||
```text
|
||||
==> 0 # Warm-up run
|
||||
[Model Preprocess]: 12.75788ms
|
||||
[ORT H2D]: 237.118µs
|
||||
[ORT Inference]: 507.895469ms
|
||||
[ORT D2H]: 191.655µs
|
||||
[Model Inference]: 508.34589ms
|
||||
[Model Postprocess]: 1.061122ms
|
||||
==> 1 # Stable run
|
||||
[Model Preprocess]: 13.658655ms
|
||||
[ORT H2D]: 209.975µs
|
||||
[ORT Inference]: 5.12372ms
|
||||
[ORT D2H]: 182.389µs
|
||||
[Model Inference]: 5.530022ms
|
||||
[Model Postprocess]: 1.04851ms
|
||||
==> 2 # Stable run
|
||||
[Model Preprocess]: 12.475332ms
|
||||
[ORT H2D]: 246.127µs
|
||||
[ORT Inference]: 5.048432ms
|
||||
[ORT D2H]: 187.117µs
|
||||
[Model Inference]: 5.493119ms
|
||||
[Model Postprocess]: 1.040906ms
|
||||
```
|
||||
|
||||
#### Other Options
|
||||
|
||||
- `--conf`: Confidence threshold for detections \[default: 0.3].
|
||||
- `--iou`: IoU (Intersection over Union) threshold for Non-Maximum Suppression (NMS) \[default: 0.45].
|
||||
- `--kconf`: Confidence threshold for keypoints (in Pose Estimation) \[default: 0.55].
|
||||
- `--plot`: Plot the inference results with random RGB colors and save the output image to the `runs` directory.
|
||||
|
||||
You can view all available command-line arguments by running:
|
||||
|
||||
```bash
|
||||
# Clone the repository if you haven't already
|
||||
# git clone https://github.com/ultralytics/ultralytics
|
||||
# cd ultralytics/examples/YOLOv8-ONNXRuntime-Rust
|
||||
|
||||
cargo run --release -- --help
|
||||
```
|
||||
|
||||
## 🖼️ Examples
|
||||
|
||||
|
||||
|
||||
### Classification
|
||||
|
||||
Running a dynamic shape ONNX classification model on the `CPU` with a specific image size (`--height 224 --width 224`). The plotted result image will be saved in the `runs` directory.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --model ../assets/weights/yolov8m-cls-dyn.onnx --source ../assets/images/dog.jpg --height 224 --width 224 --plot --profile
|
||||
```
|
||||
|
||||
Example output:
|
||||
|
||||
```text
|
||||
Summary:
|
||||
> Task: Classify (Ultralytics 8.0.217) # Version might differ
|
||||
> EP: Cpu
|
||||
> Dtype: Float32
|
||||
> Batch: 1 (Dynamic), Height: 224 (Dynamic), Width: 224 (Dynamic)
|
||||
> nc: 1000 nk: 0, nm: 0, conf: 0.3, kconf: 0.55, iou: 0.45
|
||||
|
||||
[Model Preprocess]: 16.363477ms
|
||||
[ORT H2D]: 50.722µs
|
||||
[ORT Inference]: 16.295808ms
|
||||
[ORT D2H]: 8.37µs
|
||||
[Model Inference]: 16.367046ms
|
||||
[Model Postprocess]: 3.527µs
|
||||
[
|
||||
YOLOResult {
|
||||
Probs(top5): Some([(208, 0.6950566), (209, 0.13823675), (178, 0.04849795), (215, 0.019029364), (212, 0.016506357)]), # Class IDs and confidences
|
||||
Bboxes: None,
|
||||
Keypoints: None,
|
||||
Masks: None,
|
||||
},
|
||||
]
|
||||
```
|
||||
|
||||
### Object Detection
|
||||
|
||||
Using the `CUDA` execution provider and a dynamic image size (`--height 640 --width 480`).
|
||||
|
||||
```bash
|
||||
cargo run --release -- --cuda --model ../assets/weights/yolov8m-dynamic.onnx --source ../assets/images/bus.jpg --plot --height 640 --width 480
|
||||
```
|
||||
|
||||
### Pose Detection
|
||||
|
||||
Using the `TensorRT` execution provider.
|
||||
|
||||
```bash
|
||||
cargo run --release -- --trt --model ../assets/weights/yolov8m-pose.onnx --source ../assets/images/bus.jpg --plot
|
||||
```
|
||||
|
||||
### Instance Segmentation
|
||||
|
||||
Using the `TensorRT` execution provider with an FP16 model (`--fp16`).
|
||||
|
||||
```bash
|
||||
cargo run --release -- --trt --fp16 --model ../assets/weights/yolov8m-seg.onnx --source ../assets/images/0172.jpg --plot
|
||||
```
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! If you find any issues or have suggestions for improvement, please feel free to open an issue or submit a pull request to the main [Ultralytics repository](https://github.com/ultralytics/ultralytics).
|
||||
@@ -0,0 +1,89 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
use clap::Parser;
|
||||
|
||||
use crate::YOLOTask;
|
||||
|
||||
#[derive(Parser, Clone)]
|
||||
#[command(author, version, about, long_about = None)]
|
||||
pub struct Args {
|
||||
/// ONNX model path
|
||||
#[arg(long, required = true)]
|
||||
pub model: String,
|
||||
|
||||
/// input path
|
||||
#[arg(long, required = true)]
|
||||
pub source: String,
|
||||
|
||||
/// device id
|
||||
#[arg(long, default_value_t = 0)]
|
||||
pub device_id: i32,
|
||||
|
||||
/// using TensorRT EP
|
||||
#[arg(long)]
|
||||
pub trt: bool,
|
||||
|
||||
/// using CUDA EP
|
||||
#[arg(long)]
|
||||
pub cuda: bool,
|
||||
|
||||
/// input batch size
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub batch: u32,
|
||||
|
||||
/// trt input min_batch size
|
||||
#[arg(long, default_value_t = 1)]
|
||||
pub batch_min: u32,
|
||||
|
||||
/// trt input max_batch size
|
||||
#[arg(long, default_value_t = 32)]
|
||||
pub batch_max: u32,
|
||||
|
||||
/// using TensorRT --fp16
|
||||
#[arg(long)]
|
||||
pub fp16: bool,
|
||||
|
||||
/// specify YOLO task
|
||||
#[arg(long, value_enum)]
|
||||
pub task: Option<YOLOTask>,
|
||||
|
||||
/// num_classes
|
||||
#[arg(long)]
|
||||
pub nc: Option<u32>,
|
||||
|
||||
/// num_keypoints
|
||||
#[arg(long)]
|
||||
pub nk: Option<u32>,
|
||||
|
||||
/// num_masks
|
||||
#[arg(long)]
|
||||
pub nm: Option<u32>,
|
||||
|
||||
/// input image width
|
||||
#[arg(long)]
|
||||
pub width: Option<u32>,
|
||||
|
||||
/// input image height
|
||||
#[arg(long)]
|
||||
pub height: Option<u32>,
|
||||
|
||||
/// confidence threshold
|
||||
#[arg(long, required = false, default_value_t = 0.3)]
|
||||
pub conf: f32,
|
||||
|
||||
/// iou threshold in NMS
|
||||
#[arg(long, required = false, default_value_t = 0.45)]
|
||||
pub iou: f32,
|
||||
|
||||
/// confidence threshold of keypoint
|
||||
#[arg(long, required = false, default_value_t = 0.55)]
|
||||
pub kconf: f32,
|
||||
|
||||
/// plot inference result and save
|
||||
#[arg(long)]
|
||||
pub plot: bool,
|
||||
|
||||
/// check time consumed in each stage
|
||||
#[arg(long)]
|
||||
pub profile: bool,
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
#![allow(clippy::type_complexity)]
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
use std::io::{Read, Write};
|
||||
|
||||
pub mod cli;
|
||||
pub mod model;
|
||||
pub mod ort_backend;
|
||||
pub mod yolo_result;
|
||||
pub use crate::cli::Args;
|
||||
pub use crate::model::YOLOv8;
|
||||
pub use crate::ort_backend::{Batch, OrtBackend, OrtConfig, OrtEP, YOLOTask};
|
||||
pub use crate::yolo_result::{Bbox, Embedding, Point2, YOLOResult};
|
||||
|
||||
pub fn non_max_suppression(
|
||||
xs: &mut Vec<(Bbox, Option<Vec<Point2>>, Option<Vec<f32>>)>,
|
||||
iou_threshold: f32,
|
||||
) {
|
||||
xs.sort_by(|b1, b2| b2.0.confidence().partial_cmp(&b1.0.confidence()).unwrap());
|
||||
|
||||
let mut current_index = 0;
|
||||
for index in 0..xs.len() {
|
||||
let mut drop = false;
|
||||
for prev_index in 0..current_index {
|
||||
let iou = xs[prev_index].0.iou(&xs[index].0);
|
||||
if iou > iou_threshold {
|
||||
drop = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if !drop {
|
||||
xs.swap(current_index, index);
|
||||
current_index += 1;
|
||||
}
|
||||
}
|
||||
xs.truncate(current_index);
|
||||
}
|
||||
|
||||
pub fn gen_time_string(delimiter: &str) -> String {
|
||||
let offset = chrono::FixedOffset::east_opt(8 * 60 * 60).unwrap(); // Beijing
|
||||
let t_now = chrono::Utc::now().with_timezone(&offset);
|
||||
let fmt = format!(
|
||||
"%Y{}%m{}%d{}%H{}%M{}%S{}%f",
|
||||
delimiter, delimiter, delimiter, delimiter, delimiter, delimiter
|
||||
);
|
||||
t_now.format(&fmt).to_string()
|
||||
}
|
||||
|
||||
pub const SKELETON: [(usize, usize); 16] = [
|
||||
(0, 1),
|
||||
(0, 2),
|
||||
(1, 3),
|
||||
(2, 4),
|
||||
(5, 6),
|
||||
(5, 11),
|
||||
(6, 12),
|
||||
(11, 12),
|
||||
(5, 7),
|
||||
(6, 8),
|
||||
(7, 9),
|
||||
(8, 10),
|
||||
(11, 13),
|
||||
(12, 14),
|
||||
(13, 15),
|
||||
(14, 16),
|
||||
];
|
||||
|
||||
pub fn check_font(font: &str) -> rusttype::Font<'static> {
|
||||
// check then load font
|
||||
|
||||
// ultralytics font path
|
||||
let font_path_config = match dirs::config_dir() {
|
||||
Some(mut d) => {
|
||||
d.push("Ultralytics");
|
||||
d.push(font);
|
||||
d
|
||||
}
|
||||
None => panic!("Unsupported operating system. Now support Linux, MacOS, Windows."),
|
||||
};
|
||||
|
||||
// current font path
|
||||
let font_path_current = std::path::PathBuf::from(font);
|
||||
|
||||
// check font
|
||||
let font_path = if font_path_config.exists() {
|
||||
font_path_config
|
||||
} else if font_path_current.exists() {
|
||||
font_path_current
|
||||
} else {
|
||||
println!("Downloading font...");
|
||||
let source_url = "https://ultralytics.com/assets/Arial.ttf";
|
||||
let resp = ureq::get(source_url)
|
||||
.timeout(std::time::Duration::from_secs(500))
|
||||
.call()
|
||||
.unwrap_or_else(|err| panic!("> Failed to download font: {source_url}: {err:?}"));
|
||||
|
||||
// read to buffer with size limit (10MB max for font file)
|
||||
const MAX_FONT_SIZE: u64 = 10 * 1024 * 1024;
|
||||
let mut buffer = vec![];
|
||||
resp.into_reader()
|
||||
.take(MAX_FONT_SIZE)
|
||||
.read_to_end(&mut buffer)
|
||||
.unwrap();
|
||||
|
||||
// save
|
||||
let _path = std::fs::File::create(font).unwrap();
|
||||
let mut writer = std::io::BufWriter::new(_path);
|
||||
writer.write_all(&buffer).unwrap();
|
||||
println!("Font saved at: {:?}", font_path_current.display());
|
||||
font_path_current
|
||||
};
|
||||
|
||||
// load font
|
||||
let buffer = std::fs::read(font_path).unwrap();
|
||||
rusttype::Font::try_from_vec(buffer).unwrap()
|
||||
}
|
||||
|
||||
use ab_glyph::FontArc;
|
||||
pub fn load_font() -> FontArc {
|
||||
use std::path::Path;
|
||||
let font_path = Path::new("./font/Arial.ttf");
|
||||
match font_path.try_exists() {
|
||||
Ok(true) => {
|
||||
let buffer = std::fs::read(font_path).unwrap();
|
||||
FontArc::try_from_vec(buffer).unwrap()
|
||||
}
|
||||
Ok(false) => {
|
||||
std::fs::create_dir_all("./font").unwrap();
|
||||
println!("Downloading font...");
|
||||
let source_url = "https://ultralytics.com/assets/Arial.ttf";
|
||||
let resp = ureq::get(source_url)
|
||||
.timeout(std::time::Duration::from_secs(500))
|
||||
.call()
|
||||
.unwrap_or_else(|err| panic!("> Failed to download font: {source_url}: {err:?}"));
|
||||
|
||||
// read to buffer with size limit (10MB max for font file)
|
||||
const MAX_FONT_SIZE: u64 = 10 * 1024 * 1024;
|
||||
let mut buffer = vec![];
|
||||
resp.into_reader()
|
||||
.take(MAX_FONT_SIZE)
|
||||
.read_to_end(&mut buffer)
|
||||
.unwrap();
|
||||
// save
|
||||
let mut fd = std::fs::File::create(font_path).unwrap();
|
||||
fd.write_all(&buffer).unwrap();
|
||||
println!("Font saved at: {:?}", font_path.display());
|
||||
FontArc::try_from_vec(buffer).unwrap()
|
||||
}
|
||||
Err(e) => {
|
||||
panic!("Failed to load font {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
use clap::Parser;
|
||||
|
||||
use yolov8_rs::{Args, YOLOv8};
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let args = Args::parse();
|
||||
|
||||
// 1. load image
|
||||
let x = image::ImageReader::open(&args.source)?
|
||||
.with_guessed_format()?
|
||||
.decode()?;
|
||||
|
||||
// 2. model support dynamic batch inference, so input should be a Vec
|
||||
let xs = vec![x];
|
||||
|
||||
// You can test `--batch 2` with this
|
||||
// let xs = vec![x.clone(), x];
|
||||
|
||||
// 3. build yolov8 model
|
||||
let mut model = YOLOv8::new(args)?;
|
||||
model.summary(); // model info
|
||||
|
||||
// 4. run
|
||||
let ys = model.run(&xs)?;
|
||||
println!("{:?}", ys);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,652 @@
|
||||
#![allow(clippy::type_complexity)]
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
use ab_glyph::FontArc;
|
||||
use anyhow::Result;
|
||||
use image::{DynamicImage, GenericImageView, ImageBuffer};
|
||||
use ndarray::{s, Array, Axis, IxDyn};
|
||||
use rand::{thread_rng, Rng};
|
||||
use std::path::PathBuf;
|
||||
|
||||
use crate::{
|
||||
gen_time_string, load_font, non_max_suppression, Args, Batch, Bbox, Embedding, OrtBackend,
|
||||
OrtConfig, OrtEP, Point2, YOLOResult, YOLOTask, SKELETON,
|
||||
};
|
||||
|
||||
pub struct YOLOv8 {
|
||||
// YOLOv8 model for all yolo-tasks
|
||||
engine: OrtBackend,
|
||||
nc: u32,
|
||||
nk: u32,
|
||||
nm: u32,
|
||||
height: u32,
|
||||
width: u32,
|
||||
batch: u32,
|
||||
task: YOLOTask,
|
||||
conf: f32,
|
||||
kconf: f32,
|
||||
iou: f32,
|
||||
names: Vec<String>,
|
||||
color_palette: Vec<(u8, u8, u8)>,
|
||||
profile: bool,
|
||||
plot: bool,
|
||||
}
|
||||
|
||||
impl YOLOv8 {
|
||||
pub fn new(config: Args) -> Result<Self> {
|
||||
// execution provider
|
||||
let ep = if config.trt {
|
||||
OrtEP::Trt(config.device_id)
|
||||
} else if config.cuda {
|
||||
OrtEP::CUDA(config.device_id)
|
||||
} else {
|
||||
OrtEP::CPU
|
||||
};
|
||||
|
||||
// batch
|
||||
let batch = Batch {
|
||||
opt: config.batch,
|
||||
min: config.batch_min,
|
||||
max: config.batch_max,
|
||||
};
|
||||
|
||||
// build ort engine
|
||||
let ort_args = OrtConfig {
|
||||
ep,
|
||||
batch,
|
||||
f: config.model,
|
||||
task: config.task,
|
||||
trt_fp16: config.fp16,
|
||||
image_size: (config.height, config.width),
|
||||
};
|
||||
let engine = OrtBackend::build(ort_args)?;
|
||||
|
||||
// get batch, height, width, tasks, nc, nk, nm
|
||||
let (batch, height, width, task) = (
|
||||
engine.batch(),
|
||||
engine.height(),
|
||||
engine.width(),
|
||||
engine.task(),
|
||||
);
|
||||
let nc = engine.nc().or(config.nc).unwrap_or_else(|| {
|
||||
panic!("Failed to get num_classes, make it explicit with `--nc`");
|
||||
});
|
||||
let (nk, nm) = match task {
|
||||
YOLOTask::Pose => {
|
||||
let nk = engine.nk().or(config.nk).unwrap_or_else(|| {
|
||||
panic!("Failed to get num_keypoints, make it explicit with `--nk`");
|
||||
});
|
||||
(nk, 0)
|
||||
}
|
||||
YOLOTask::Segment => {
|
||||
let nm = engine.nm().or(config.nm).unwrap_or_else(|| {
|
||||
panic!("Failed to get num_masks, make it explicit with `--nm`");
|
||||
});
|
||||
(0, nm)
|
||||
}
|
||||
_ => (0, 0),
|
||||
};
|
||||
|
||||
// class names
|
||||
let names = engine.names().unwrap_or(vec!["Unknown".to_string()]);
|
||||
|
||||
// color palette
|
||||
let mut rng = thread_rng();
|
||||
let color_palette: Vec<_> = names
|
||||
.iter()
|
||||
.map(|_| {
|
||||
(
|
||||
rng.gen_range(0..=255),
|
||||
rng.gen_range(0..=255),
|
||||
rng.gen_range(0..=255),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Self {
|
||||
engine,
|
||||
names,
|
||||
conf: config.conf,
|
||||
kconf: config.kconf,
|
||||
iou: config.iou,
|
||||
color_palette,
|
||||
profile: config.profile,
|
||||
plot: config.plot,
|
||||
nc,
|
||||
nk,
|
||||
nm,
|
||||
height,
|
||||
width,
|
||||
batch,
|
||||
task,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn scale_wh(&self, w0: f32, h0: f32, w1: f32, h1: f32) -> (f32, f32, f32) {
|
||||
let r = (w1 / w0).min(h1 / h0);
|
||||
(r, (w0 * r).round(), (h0 * r).round())
|
||||
}
|
||||
|
||||
pub fn preprocess(&mut self, xs: &Vec<DynamicImage>) -> Result<Array<f32, IxDyn>> {
|
||||
let mut ys =
|
||||
Array::ones((xs.len(), 3, self.height() as usize, self.width() as usize)).into_dyn();
|
||||
ys.fill(144.0 / 255.0);
|
||||
for (idx, x) in xs.iter().enumerate() {
|
||||
let img = match self.task() {
|
||||
YOLOTask::Classify => x.resize_exact(
|
||||
self.width(),
|
||||
self.height(),
|
||||
image::imageops::FilterType::Triangle,
|
||||
),
|
||||
_ => {
|
||||
let (w0, h0) = x.dimensions();
|
||||
let w0 = w0 as f32;
|
||||
let h0 = h0 as f32;
|
||||
let (_, w_new, h_new) =
|
||||
self.scale_wh(w0, h0, self.width() as f32, self.height() as f32); // f32 round
|
||||
x.resize_exact(
|
||||
w_new as u32,
|
||||
h_new as u32,
|
||||
if let YOLOTask::Segment = self.task() {
|
||||
image::imageops::FilterType::CatmullRom
|
||||
} else {
|
||||
image::imageops::FilterType::Triangle
|
||||
},
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
for (x, y, rgb) in img.pixels() {
|
||||
let x = x as usize;
|
||||
let y = y as usize;
|
||||
let [r, g, b, _] = rgb.0;
|
||||
ys[[idx, 0, y, x]] = (r as f32) / 255.0;
|
||||
ys[[idx, 1, y, x]] = (g as f32) / 255.0;
|
||||
ys[[idx, 2, y, x]] = (b as f32) / 255.0;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ys)
|
||||
}
|
||||
|
||||
pub fn run(&mut self, xs: &Vec<DynamicImage>) -> Result<Vec<YOLOResult>> {
|
||||
// pre-process
|
||||
let t_pre = std::time::Instant::now();
|
||||
let xs_ = self.preprocess(xs)?;
|
||||
if self.profile {
|
||||
println!("[Model Preprocess]: {:?}", t_pre.elapsed());
|
||||
}
|
||||
|
||||
// run
|
||||
let t_run = std::time::Instant::now();
|
||||
let ys = self.engine.run(xs_, self.profile)?;
|
||||
if self.profile {
|
||||
println!("[Model Inference]: {:?}", t_run.elapsed());
|
||||
}
|
||||
|
||||
// post-process
|
||||
let t_post = std::time::Instant::now();
|
||||
let ys = self.postprocess(ys, xs)?;
|
||||
if self.profile {
|
||||
println!("[Model Postprocess]: {:?}", t_post.elapsed());
|
||||
}
|
||||
|
||||
// plot and save
|
||||
if self.plot {
|
||||
self.plot_and_save(&ys, xs, Some(&SKELETON));
|
||||
}
|
||||
Ok(ys)
|
||||
}
|
||||
|
||||
pub fn postprocess(
|
||||
&self,
|
||||
xs: Vec<Array<f32, IxDyn>>,
|
||||
xs0: &[DynamicImage],
|
||||
) -> Result<Vec<YOLOResult>> {
|
||||
if let YOLOTask::Classify = self.task() {
|
||||
let mut ys = Vec::new();
|
||||
let preds = &xs[0];
|
||||
for batch in preds.axis_iter(Axis(0)) {
|
||||
ys.push(YOLOResult::new(
|
||||
Some(Embedding::new(batch.into_owned())),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
));
|
||||
}
|
||||
Ok(ys)
|
||||
} else {
|
||||
const CXYWH_OFFSET: usize = 4; // cxcywh
|
||||
const KPT_STEP: usize = 3; // xyconf
|
||||
let preds = &xs[0];
|
||||
let protos = {
|
||||
if xs.len() > 1 {
|
||||
Some(&xs[1])
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
let mut ys = Vec::new();
|
||||
for (idx, anchor) in preds.axis_iter(Axis(0)).enumerate() {
|
||||
// [bs, 4 + nc + nm, anchors]
|
||||
// input image
|
||||
let width_original = xs0[idx].width() as f32;
|
||||
let height_original = xs0[idx].height() as f32;
|
||||
let ratio = (self.width() as f32 / width_original)
|
||||
.min(self.height() as f32 / height_original);
|
||||
|
||||
// save each result
|
||||
let mut data: Vec<(Bbox, Option<Vec<Point2>>, Option<Vec<f32>>)> = Vec::new();
|
||||
for pred in anchor.axis_iter(Axis(1)) {
|
||||
// split preds for different tasks
|
||||
let bbox = pred.slice(s![0..CXYWH_OFFSET]);
|
||||
let clss = pred.slice(s![CXYWH_OFFSET..CXYWH_OFFSET + self.nc() as usize]);
|
||||
let kpts = {
|
||||
if let YOLOTask::Pose = self.task() {
|
||||
Some(pred.slice(s![pred.len() - KPT_STEP * self.nk() as usize..]))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
let coefs = {
|
||||
if let YOLOTask::Segment = self.task() {
|
||||
Some(pred.slice(s![pred.len() - self.nm() as usize..]).to_vec())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// confidence and id
|
||||
let (id, &confidence) = clss
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.reduce(|max, x| if x.1 > max.1 { x } else { max })
|
||||
.unwrap(); // definitely will not panic!
|
||||
|
||||
// confidence filter
|
||||
if confidence < self.conf {
|
||||
continue;
|
||||
}
|
||||
|
||||
// bbox re-scale
|
||||
let cx = bbox[0] / ratio;
|
||||
let cy = bbox[1] / ratio;
|
||||
let w = bbox[2] / ratio;
|
||||
let h = bbox[3] / ratio;
|
||||
let x = cx - w / 2.;
|
||||
let y = cy - h / 2.;
|
||||
let y_bbox = Bbox::new(
|
||||
x.max(0.0f32).min(width_original),
|
||||
y.max(0.0f32).min(height_original),
|
||||
w,
|
||||
h,
|
||||
id,
|
||||
confidence,
|
||||
);
|
||||
|
||||
// kpts
|
||||
let y_kpts = {
|
||||
if let Some(kpts) = kpts {
|
||||
let mut kpts_ = Vec::new();
|
||||
// rescale
|
||||
for i in 0..self.nk() as usize {
|
||||
let kx = kpts[KPT_STEP * i] / ratio;
|
||||
let ky = kpts[KPT_STEP * i + 1] / ratio;
|
||||
let kconf = kpts[KPT_STEP * i + 2];
|
||||
if kconf < self.kconf {
|
||||
kpts_.push(Point2::default());
|
||||
} else {
|
||||
kpts_.push(Point2::new_with_conf(
|
||||
kx.max(0.0f32).min(width_original),
|
||||
ky.max(0.0f32).min(height_original),
|
||||
kconf,
|
||||
));
|
||||
}
|
||||
}
|
||||
Some(kpts_)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// data merged
|
||||
data.push((y_bbox, y_kpts, coefs));
|
||||
}
|
||||
|
||||
// nms
|
||||
non_max_suppression(&mut data, self.iou);
|
||||
|
||||
// decode
|
||||
let mut y_bboxes: Vec<Bbox> = Vec::new();
|
||||
let mut y_kpts: Vec<Vec<Point2>> = Vec::new();
|
||||
let mut y_masks: Vec<Vec<u8>> = Vec::new();
|
||||
for elem in data.into_iter() {
|
||||
if let Some(kpts) = elem.1 {
|
||||
y_kpts.push(kpts)
|
||||
}
|
||||
|
||||
// decode masks
|
||||
if let Some(coefs) = elem.2 {
|
||||
let proto = protos.unwrap().slice(s![idx, .., .., ..]);
|
||||
let (nm, nh, nw) = proto.dim();
|
||||
|
||||
// coefs * proto -> mask
|
||||
let coefs = Array::from_shape_vec((1, nm), coefs)?; // (n, nm)
|
||||
|
||||
let proto = proto.to_owned();
|
||||
let proto = proto.to_shape((nm, nh * nw))?; // (nm, nh*nw)
|
||||
let mask = coefs.dot(&proto); // (nh, nw, n)
|
||||
let mask = mask.to_shape((nh, nw, 1))?;
|
||||
|
||||
// build image from ndarray
|
||||
let mask_im: ImageBuffer<image::Luma<_>, Vec<f32>> =
|
||||
match ImageBuffer::from_raw(
|
||||
nw as u32,
|
||||
nh as u32,
|
||||
mask.to_owned().into_raw_vec_and_offset().0,
|
||||
) {
|
||||
Some(image) => image,
|
||||
None => panic!("can not create image from ndarray"),
|
||||
};
|
||||
let mut mask_im = image::DynamicImage::from(mask_im); // -> dyn
|
||||
|
||||
// rescale masks
|
||||
let (_, w_mask, h_mask) =
|
||||
self.scale_wh(width_original, height_original, nw as f32, nh as f32);
|
||||
let mask_cropped = mask_im.crop(0, 0, w_mask as u32, h_mask as u32);
|
||||
let mask_original = mask_cropped.resize_exact(
|
||||
// resize_to_fill
|
||||
width_original as u32,
|
||||
height_original as u32,
|
||||
match self.task() {
|
||||
YOLOTask::Segment => image::imageops::FilterType::CatmullRom,
|
||||
_ => image::imageops::FilterType::Triangle,
|
||||
},
|
||||
);
|
||||
|
||||
// crop-mask with bbox
|
||||
let mut mask_original_cropped = mask_original.into_luma8();
|
||||
for y in 0..height_original as usize {
|
||||
for x in 0..width_original as usize {
|
||||
if x < elem.0.xmin() as usize
|
||||
|| x > elem.0.xmax() as usize
|
||||
|| y < elem.0.ymin() as usize
|
||||
|| y > elem.0.ymax() as usize
|
||||
{
|
||||
mask_original_cropped.put_pixel(
|
||||
x as u32,
|
||||
y as u32,
|
||||
image::Luma([0u8]),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
y_masks.push(mask_original_cropped.into_raw());
|
||||
}
|
||||
y_bboxes.push(elem.0);
|
||||
}
|
||||
|
||||
// save each result
|
||||
let y = YOLOResult {
|
||||
probs: None,
|
||||
bboxes: if !y_bboxes.is_empty() {
|
||||
Some(y_bboxes)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
keypoints: if !y_kpts.is_empty() {
|
||||
Some(y_kpts)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
masks: if !y_masks.is_empty() {
|
||||
Some(y_masks)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
};
|
||||
ys.push(y);
|
||||
}
|
||||
|
||||
Ok(ys)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn plot_and_save(
|
||||
&self,
|
||||
ys: &[YOLOResult],
|
||||
xs0: &[DynamicImage],
|
||||
skeletons: Option<&[(usize, usize)]>,
|
||||
) {
|
||||
// check font then load
|
||||
let font: FontArc = load_font();
|
||||
for (_idb, (img0, y)) in xs0.iter().zip(ys.iter()).enumerate() {
|
||||
let mut img = img0.to_rgb8();
|
||||
|
||||
// draw for classifier
|
||||
if let Some(probs) = y.probs() {
|
||||
for (i, k) in probs.topk(5).iter().enumerate() {
|
||||
let legend = format!("{} {:.2}%", self.names[k.0], k.1);
|
||||
let scale = 32;
|
||||
let legend_size = img.width().max(img.height()) / scale;
|
||||
let x = img.width() / 20;
|
||||
let y = img.height() / 20 + i as u32 * legend_size;
|
||||
|
||||
imageproc::drawing::draw_text_mut(
|
||||
&mut img,
|
||||
image::Rgb([0, 255, 0]),
|
||||
x as i32,
|
||||
y as i32,
|
||||
legend_size as f32,
|
||||
&font,
|
||||
&legend,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// draw bboxes & keypoints
|
||||
if let Some(bboxes) = y.bboxes() {
|
||||
for (_idx, bbox) in bboxes.iter().enumerate() {
|
||||
// rect
|
||||
imageproc::drawing::draw_hollow_rect_mut(
|
||||
&mut img,
|
||||
imageproc::rect::Rect::at(bbox.xmin() as i32, bbox.ymin() as i32)
|
||||
.of_size(bbox.width() as u32, bbox.height() as u32),
|
||||
image::Rgb(self.color_palette[bbox.id()].into()),
|
||||
);
|
||||
|
||||
// text
|
||||
let legend = format!("{} {:.2}%", self.names[bbox.id()], bbox.confidence());
|
||||
let scale = 40;
|
||||
let legend_size = img.width().max(img.height()) / scale;
|
||||
imageproc::drawing::draw_text_mut(
|
||||
&mut img,
|
||||
image::Rgb(self.color_palette[bbox.id()].into()),
|
||||
bbox.xmin() as i32,
|
||||
(bbox.ymin() - legend_size as f32) as i32,
|
||||
legend_size as f32,
|
||||
&font,
|
||||
&legend,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// draw kpts
|
||||
if let Some(keypoints) = y.keypoints() {
|
||||
for kpts in keypoints.iter() {
|
||||
for kpt in kpts.iter() {
|
||||
// filter
|
||||
if kpt.confidence() < self.kconf {
|
||||
continue;
|
||||
}
|
||||
|
||||
// draw point
|
||||
imageproc::drawing::draw_filled_circle_mut(
|
||||
&mut img,
|
||||
(kpt.x() as i32, kpt.y() as i32),
|
||||
2,
|
||||
image::Rgb([0, 255, 0]),
|
||||
);
|
||||
}
|
||||
|
||||
// draw skeleton if has
|
||||
if let Some(skeletons) = skeletons {
|
||||
for &(idx1, idx2) in skeletons.iter() {
|
||||
let kpt1 = &kpts[idx1];
|
||||
let kpt2 = &kpts[idx2];
|
||||
if kpt1.confidence() < self.kconf || kpt2.confidence() < self.kconf {
|
||||
continue;
|
||||
}
|
||||
imageproc::drawing::draw_line_segment_mut(
|
||||
&mut img,
|
||||
(kpt1.x(), kpt1.y()),
|
||||
(kpt2.x(), kpt2.y()),
|
||||
image::Rgb([233, 14, 57]),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// draw mask
|
||||
if let Some(masks) = y.masks() {
|
||||
for (mask, _bbox) in masks.iter().zip(y.bboxes().unwrap().iter()) {
|
||||
let mask_nd: ImageBuffer<image::Luma<_>, Vec<u8>> =
|
||||
match ImageBuffer::from_vec(img.width(), img.height(), mask.to_vec()) {
|
||||
Some(image) => image,
|
||||
None => panic!("can not create image from ndarray"),
|
||||
};
|
||||
|
||||
for _x in 0..img.width() {
|
||||
for _y in 0..img.height() {
|
||||
let mask_p = imageproc::drawing::Canvas::get_pixel(&mask_nd, _x, _y);
|
||||
if mask_p.0[0] > 0 {
|
||||
let mut img_p = imageproc::drawing::Canvas::get_pixel(&img, _x, _y);
|
||||
// img_p.0[2] = self.color_palette[bbox.id()].2 / 2;
|
||||
// img_p.0[1] = self.color_palette[bbox.id()].1 / 2;
|
||||
// img_p.0[0] = self.color_palette[bbox.id()].0 / 2;
|
||||
img_p.0[2] /= 2;
|
||||
img_p.0[1] = 255 - (255 - img_p.0[2]) / 2;
|
||||
img_p.0[0] /= 2;
|
||||
imageproc::drawing::Canvas::draw_pixel(&mut img, _x, _y, img_p)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// mkdir and save
|
||||
let mut runs = PathBuf::from("runs");
|
||||
if !runs.exists() {
|
||||
std::fs::create_dir_all(&runs).unwrap();
|
||||
}
|
||||
runs.push(gen_time_string("-"));
|
||||
let saveout = format!("{}.jpg", runs.to_str().unwrap());
|
||||
let _ = img.save(saveout);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn summary(&self) {
|
||||
println!(
|
||||
"\nSummary:\n\
|
||||
> Task: {:?}{}\n\
|
||||
> EP: {:?} {}\n\
|
||||
> Dtype: {:?}\n\
|
||||
> Batch: {} ({}), Height: {} ({}), Width: {} ({})\n\
|
||||
> nc: {} nk: {}, nm: {}, conf: {}, kconf: {}, iou: {}\n\
|
||||
",
|
||||
self.task(),
|
||||
match self.engine.author().zip(self.engine.version()) {
|
||||
Some((author, ver)) => format!(" ({} {})", author, ver),
|
||||
None => String::from(""),
|
||||
},
|
||||
self.engine.ep(),
|
||||
if let OrtEP::CPU = self.engine.ep() {
|
||||
""
|
||||
} else {
|
||||
"(May still fall back to CPU)"
|
||||
},
|
||||
self.engine.dtype(),
|
||||
self.batch(),
|
||||
if self.engine.is_batch_dynamic() {
|
||||
"Dynamic"
|
||||
} else {
|
||||
"Const"
|
||||
},
|
||||
self.height(),
|
||||
if self.engine.is_height_dynamic() {
|
||||
"Dynamic"
|
||||
} else {
|
||||
"Const"
|
||||
},
|
||||
self.width(),
|
||||
if self.engine.is_width_dynamic() {
|
||||
"Dynamic"
|
||||
} else {
|
||||
"Const"
|
||||
},
|
||||
self.nc(),
|
||||
self.nk(),
|
||||
self.nm(),
|
||||
self.conf,
|
||||
self.kconf,
|
||||
self.iou,
|
||||
);
|
||||
}
|
||||
|
||||
pub fn engine(&self) -> &OrtBackend {
|
||||
&self.engine
|
||||
}
|
||||
|
||||
pub fn conf(&self) -> f32 {
|
||||
self.conf
|
||||
}
|
||||
|
||||
pub fn set_conf(&mut self, val: f32) {
|
||||
self.conf = val;
|
||||
}
|
||||
|
||||
pub fn conf_mut(&mut self) -> &mut f32 {
|
||||
&mut self.conf
|
||||
}
|
||||
|
||||
pub fn kconf(&self) -> f32 {
|
||||
self.kconf
|
||||
}
|
||||
|
||||
pub fn iou(&self) -> f32 {
|
||||
self.iou
|
||||
}
|
||||
|
||||
pub fn task(&self) -> &YOLOTask {
|
||||
&self.task
|
||||
}
|
||||
|
||||
pub fn batch(&self) -> u32 {
|
||||
self.batch
|
||||
}
|
||||
|
||||
pub fn width(&self) -> u32 {
|
||||
self.width
|
||||
}
|
||||
|
||||
pub fn height(&self) -> u32 {
|
||||
self.height
|
||||
}
|
||||
|
||||
pub fn nc(&self) -> u32 {
|
||||
self.nc
|
||||
}
|
||||
|
||||
pub fn nk(&self) -> u32 {
|
||||
self.nk
|
||||
}
|
||||
|
||||
pub fn nm(&self) -> u32 {
|
||||
self.nm
|
||||
}
|
||||
|
||||
pub fn names(&self) -> &Vec<String> {
|
||||
&self.names
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,609 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
use anyhow::Result;
|
||||
use clap::ValueEnum;
|
||||
use half::f16;
|
||||
use ndarray::{Array, CowArray, IxDyn};
|
||||
use ort::execution_providers::{
|
||||
CPUExecutionProvider, CUDAExecutionProvider, ExecutionProvider, ExecutionProviderDispatch,
|
||||
TensorRTExecutionProvider,
|
||||
};
|
||||
use ort::session::builder::SessionBuilder;
|
||||
use ort::session::Session;
|
||||
use ort::tensor::TensorElementType;
|
||||
use ort::value::ValueType;
|
||||
use regex::Regex;
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
|
||||
pub enum YOLOTask {
|
||||
// YOLO tasks
|
||||
Classify,
|
||||
Detect,
|
||||
Pose,
|
||||
Segment,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum OrtEP {
|
||||
// ONNXRuntime execution provider
|
||||
CPU,
|
||||
CUDA(i32),
|
||||
Trt(i32),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Batch {
|
||||
pub opt: u32,
|
||||
pub min: u32,
|
||||
pub max: u32,
|
||||
}
|
||||
|
||||
impl Default for Batch {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
opt: 1,
|
||||
min: 1,
|
||||
max: 1,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct OrtInputs {
|
||||
// ONNX model inputs attrs
|
||||
pub shapes: Vec<Vec<i64>>,
|
||||
//pub dtypes: Vec<TensorElementDataType>,
|
||||
pub dtypes: Vec<TensorElementType>,
|
||||
pub names: Vec<String>,
|
||||
pub sizes: Vec<Vec<u32>>,
|
||||
}
|
||||
|
||||
impl OrtInputs {
|
||||
pub fn new(session: &Session) -> Self {
|
||||
let mut shapes = Vec::new();
|
||||
let mut dtypes = Vec::new();
|
||||
let mut names = Vec::new();
|
||||
for i in session.inputs.iter() {
|
||||
/* let shape: Vec<i32> = i
|
||||
.dimensions()
|
||||
.map(|x| if let Some(x) = x { x as i32 } else { -1i32 })
|
||||
.collect();
|
||||
shapes.push(shape); */
|
||||
if let ValueType::Tensor {
|
||||
ty,
|
||||
dimension_symbols: _,
|
||||
shape,
|
||||
} = &i.input_type
|
||||
{
|
||||
dtypes.push(ty.clone());
|
||||
let shape = shape.to_vec().clone();
|
||||
shapes.push(shape);
|
||||
} else {
|
||||
panic!("不支持的数据格式, {} - {}", file!(), line!());
|
||||
}
|
||||
//dtypes.push(i.input_type);
|
||||
names.push(i.name.clone());
|
||||
}
|
||||
Self {
|
||||
shapes,
|
||||
dtypes,
|
||||
names,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct OrtConfig {
|
||||
// ORT config
|
||||
pub f: String,
|
||||
pub task: Option<YOLOTask>,
|
||||
pub ep: OrtEP,
|
||||
pub trt_fp16: bool,
|
||||
pub batch: Batch,
|
||||
pub image_size: (Option<u32>, Option<u32>),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct OrtBackend {
|
||||
// ORT engine
|
||||
session: Session,
|
||||
task: YOLOTask,
|
||||
ep: OrtEP,
|
||||
batch: Batch,
|
||||
inputs: OrtInputs,
|
||||
}
|
||||
|
||||
impl OrtBackend {
|
||||
pub fn build(args: OrtConfig) -> Result<Self> {
|
||||
// build env & session
|
||||
// in version 2.x environment is removed
|
||||
/* let env = ort::EnvironmentBuilder
|
||||
::with_name("YOLOv8")
|
||||
.build()?
|
||||
.into_arc(); */
|
||||
let sessionbuilder = SessionBuilder::new()?;
|
||||
let session = sessionbuilder.commit_from_file(&args.f)?;
|
||||
//let session = SessionBuilder::new(&env)?.with_model_from_file(&args.f)?;
|
||||
|
||||
// get inputs
|
||||
let mut inputs = OrtInputs::new(&session);
|
||||
|
||||
// batch size
|
||||
let mut batch = args.batch;
|
||||
let batch = if inputs.shapes[0][0] == -1 {
|
||||
batch
|
||||
} else {
|
||||
assert_eq!(
|
||||
inputs.shapes[0][0] as u32, batch.opt,
|
||||
"Expected batch size: {}, got {}. Try using `--batch {}`.",
|
||||
inputs.shapes[0][0] as u32, batch.opt, inputs.shapes[0][0] as u32
|
||||
);
|
||||
batch.opt = inputs.shapes[0][0] as u32;
|
||||
batch
|
||||
};
|
||||
|
||||
// input size: height and width
|
||||
let height = if inputs.shapes[0][2] == -1 {
|
||||
match args.image_size.0 {
|
||||
Some(height) => height,
|
||||
None => panic!("Failed to get model height. Make it explicit with `--height`"),
|
||||
}
|
||||
} else {
|
||||
inputs.shapes[0][2] as u32
|
||||
};
|
||||
let width = if inputs.shapes[0][3] == -1 {
|
||||
match args.image_size.1 {
|
||||
Some(width) => width,
|
||||
None => panic!("Failed to get model width. Make it explicit with `--width`"),
|
||||
}
|
||||
} else {
|
||||
inputs.shapes[0][3] as u32
|
||||
};
|
||||
inputs.sizes.push(vec![height, width]);
|
||||
|
||||
// build provider
|
||||
let (ep, provider) = match args.ep {
|
||||
OrtEP::CUDA(device_id) => Self::set_ep_cuda(device_id),
|
||||
OrtEP::Trt(device_id) => Self::set_ep_trt(device_id, args.trt_fp16, &batch, &inputs),
|
||||
_ => (
|
||||
OrtEP::CPU,
|
||||
ExecutionProviderDispatch::from(CPUExecutionProvider::default()),
|
||||
),
|
||||
};
|
||||
|
||||
// build session again with the new provider
|
||||
let session = SessionBuilder::new()?
|
||||
// .with_optimization_level(ort::GraphOptimizationLevel::Level3)?
|
||||
.with_execution_providers([provider])?
|
||||
.commit_from_file(args.f)?;
|
||||
|
||||
// task: using given one or guessing
|
||||
let task = match args.task {
|
||||
Some(task) => task,
|
||||
None => match session.metadata() {
|
||||
Err(_) => panic!("No metadata found. Try making it explicit by `--task`"),
|
||||
Ok(metadata) => match metadata.custom("task") {
|
||||
Err(_) => panic!("Can not get custom value. Try making it explicit by `--task`"),
|
||||
Ok(value) => match value {
|
||||
None => panic!("No corresponding value of `task` found in metadata. Make it explicit by `--task`"),
|
||||
Some(task) => match task.as_str() {
|
||||
"classify" => YOLOTask::Classify,
|
||||
"detect" => YOLOTask::Detect,
|
||||
"pose" => YOLOTask::Pose,
|
||||
"segment" => YOLOTask::Segment,
|
||||
x => todo!("{:?} is not supported for now!", x),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
session,
|
||||
task,
|
||||
ep,
|
||||
batch,
|
||||
inputs,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn fetch_inputs_from_session(
|
||||
session: &Session,
|
||||
) -> (Vec<Vec<i64>>, Vec<TensorElementType>, Vec<String>) {
|
||||
// get inputs attrs from ONNX model
|
||||
let mut shapes = Vec::new();
|
||||
let mut dtypes = Vec::new();
|
||||
let mut names = Vec::new();
|
||||
for i in session.inputs.iter() {
|
||||
if let ValueType::Tensor {
|
||||
ty,
|
||||
dimension_symbols: _,
|
||||
shape,
|
||||
} = &i.input_type
|
||||
{
|
||||
dtypes.push(ty.clone());
|
||||
let shape = shape.to_vec().clone();
|
||||
shapes.push(shape);
|
||||
} else {
|
||||
panic!("不支持的数据格式, {} - {}", file!(), line!());
|
||||
}
|
||||
names.push(i.name.clone());
|
||||
}
|
||||
(shapes, dtypes, names)
|
||||
}
|
||||
|
||||
pub fn set_ep_cuda(device_id: i32) -> (OrtEP, ExecutionProviderDispatch) {
|
||||
let cuda_provider = CUDAExecutionProvider::default().with_device_id(device_id);
|
||||
if let Ok(true) = cuda_provider.is_available() {
|
||||
(
|
||||
OrtEP::CUDA(device_id),
|
||||
ExecutionProviderDispatch::from(cuda_provider), //PlantForm::CUDA(cuda_provider)
|
||||
)
|
||||
} else {
|
||||
println!("> CUDA is not available! Using CPU.");
|
||||
(
|
||||
OrtEP::CPU,
|
||||
ExecutionProviderDispatch::from(CPUExecutionProvider::default()), //PlantForm::CPU(CPUExecutionProvider::default())
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_ep_trt(
|
||||
device_id: i32,
|
||||
fp16: bool,
|
||||
batch: &Batch,
|
||||
inputs: &OrtInputs,
|
||||
) -> (OrtEP, ExecutionProviderDispatch) {
|
||||
// set TensorRT
|
||||
let trt_provider = TensorRTExecutionProvider::default().with_device_id(device_id);
|
||||
|
||||
//trt_provider.
|
||||
if let Ok(true) = trt_provider.is_available() {
|
||||
let (height, width) = (inputs.sizes[0][0], inputs.sizes[0][1]);
|
||||
if inputs.dtypes[0] == TensorElementType::Float16 && !fp16 {
|
||||
panic!(
|
||||
"Dtype mismatch! Expected: Float32, got: {:?}. You should use `--fp16`",
|
||||
inputs.dtypes[0]
|
||||
);
|
||||
}
|
||||
// dynamic shape: input_tensor_1:dim_1xdim_2x...,input_tensor_2:dim_3xdim_4x...,...
|
||||
let mut opt_string = String::new();
|
||||
let mut min_string = String::new();
|
||||
let mut max_string = String::new();
|
||||
for name in inputs.names.iter() {
|
||||
let s_opt = format!("{}:{}x3x{}x{},", name, batch.opt, height, width);
|
||||
let s_min = format!("{}:{}x3x{}x{},", name, batch.min, height, width);
|
||||
let s_max = format!("{}:{}x3x{}x{},", name, batch.max, height, width);
|
||||
opt_string.push_str(s_opt.as_str());
|
||||
min_string.push_str(s_min.as_str());
|
||||
max_string.push_str(s_max.as_str());
|
||||
}
|
||||
let _ = opt_string.pop();
|
||||
let _ = min_string.pop();
|
||||
let _ = max_string.pop();
|
||||
|
||||
let trt_provider = trt_provider
|
||||
.with_profile_opt_shapes(opt_string)
|
||||
.with_profile_min_shapes(min_string)
|
||||
.with_profile_max_shapes(max_string)
|
||||
.with_fp16(fp16)
|
||||
.with_timing_cache(true);
|
||||
(
|
||||
OrtEP::Trt(device_id),
|
||||
ExecutionProviderDispatch::from(trt_provider),
|
||||
)
|
||||
} else {
|
||||
println!("> TensorRT is not available! Try using CUDA...");
|
||||
Self::set_ep_cuda(device_id)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn fetch_from_metadata(&self, key: &str) -> Option<String> {
|
||||
// fetch value from onnx model file by key
|
||||
match self.session.metadata() {
|
||||
Err(_) => None,
|
||||
Ok(metadata) => match metadata.custom(key) {
|
||||
Err(_) => None,
|
||||
Ok(value) => value,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn run(&mut self, xs: Array<f32, IxDyn>, profile: bool) -> Result<Vec<Array<f32, IxDyn>>> {
|
||||
// ORT inference
|
||||
match self.dtype() {
|
||||
TensorElementType::Float16 => self.run_fp16(xs, profile),
|
||||
TensorElementType::Float32 => self.run_fp32(xs, profile),
|
||||
_ => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn run_fp16(
|
||||
&mut self,
|
||||
xs: Array<f32, IxDyn>,
|
||||
profile: bool,
|
||||
) -> Result<Vec<Array<f32, IxDyn>>> {
|
||||
// f32->f16
|
||||
let t = std::time::Instant::now();
|
||||
let xs = xs.mapv(f16::from_f32);
|
||||
if profile {
|
||||
println!("[ORT f32->f16]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// h2d
|
||||
let t = std::time::Instant::now();
|
||||
let xs = CowArray::from(xs);
|
||||
if profile {
|
||||
println!("[ORT H2D]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// prepare input Value from the ndarray (needed because SessionInputValue implements From<Value<_>>)
|
||||
let t = std::time::Instant::now();
|
||||
let input = ort::value::Value::from_array(xs.into_owned())?;
|
||||
if profile {
|
||||
println!("[ORT Prepare Value]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// compute output shapes before calling session.run to avoid borrowing self immutably while session is mutably borrowed
|
||||
let out_shapes = self.output_shapes();
|
||||
|
||||
// run
|
||||
let t = std::time::Instant::now();
|
||||
let ys = self.session.run(ort::inputs![input])?;
|
||||
if profile {
|
||||
println!("[ORT Inference]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// d2h
|
||||
Ok(ys
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, (_k, v))| {
|
||||
// d2h
|
||||
let t = std::time::Instant::now();
|
||||
// try_extract_tensor for f16 returns (shape, slice)
|
||||
let (_shape, slice) = v.try_extract_tensor::<f16>().unwrap();
|
||||
if profile {
|
||||
println!("[ORT D2H]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// f16->f32
|
||||
let t_ = std::time::Instant::now();
|
||||
// build ndarray from the returned slice using the runtime output shape
|
||||
let out_shape = out_shapes[idx].clone();
|
||||
let dims = out_shape.iter().map(|&d| d as usize).collect::<Vec<_>>();
|
||||
let arr_f16 = Array::from_shape_vec(IxDyn(&dims), slice.to_vec()).unwrap();
|
||||
let v = arr_f16.mapv(f16::to_f32);
|
||||
if profile {
|
||||
println!("[ORT f16->f32]: {:?}", t_.elapsed());
|
||||
}
|
||||
v
|
||||
})
|
||||
.collect::<Vec<Array<_, _>>>())
|
||||
}
|
||||
|
||||
pub fn run_fp32(
|
||||
&mut self,
|
||||
xs: Array<f32, IxDyn>,
|
||||
profile: bool,
|
||||
) -> Result<Vec<Array<f32, IxDyn>>> {
|
||||
// h2d
|
||||
let t = std::time::Instant::now();
|
||||
let xs = CowArray::from(xs);
|
||||
if profile {
|
||||
println!("[ORT H2D]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// prepare input Value from the ndarray (needed because SessionInputValue implements From<Value<_>>)
|
||||
let t = std::time::Instant::now();
|
||||
let input = ort::value::Value::from_array(xs.into_owned())?;
|
||||
if profile {
|
||||
println!("[ORT Prepare Value]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// compute output shapes before calling session.run to avoid borrowing self immutably while session is mutably borrowed
|
||||
let out_shapes = self.output_shapes();
|
||||
|
||||
// run
|
||||
let t = std::time::Instant::now();
|
||||
let ys = self.session.run(ort::inputs![input])?;
|
||||
if profile {
|
||||
println!("[ORT Inference]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// d2h
|
||||
Ok(ys
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, (_k, v))| {
|
||||
let t = std::time::Instant::now();
|
||||
// try_extract_tensor for f32 returns (shape, slice)
|
||||
let (_shape, slice) = v.try_extract_tensor::<f32>().unwrap();
|
||||
if profile {
|
||||
println!("[ORT D2H]: {:?}", t.elapsed());
|
||||
}
|
||||
|
||||
// build ndarray from the returned slice using the runtime output shape
|
||||
let out_shape = out_shapes[idx].clone();
|
||||
let dims = out_shape.iter().map(|&d| d as usize).collect::<Vec<_>>();
|
||||
Array::from_shape_vec(IxDyn(&dims), slice.to_vec()).unwrap()
|
||||
})
|
||||
.collect::<Vec<Array<f32, IxDyn>>>())
|
||||
}
|
||||
|
||||
pub fn output_shapes(&self) -> Vec<Vec<i64>> {
|
||||
let mut shapes = Vec::new();
|
||||
for output in &self.session.outputs {
|
||||
if let ValueType::Tensor { shape, .. } = &output.output_type {
|
||||
shapes.push(shape.to_vec().clone());
|
||||
} else {
|
||||
panic!("not support data format, {} - {}", file!(), line!());
|
||||
}
|
||||
}
|
||||
shapes
|
||||
}
|
||||
|
||||
pub fn output_dtypes(&self) -> Vec<TensorElementType> {
|
||||
let mut dtypes = Vec::new();
|
||||
for output in &self.session.outputs {
|
||||
if let ValueType::Tensor {
|
||||
ty,
|
||||
shape: _,
|
||||
dimension_symbols: _,
|
||||
} = &output.output_type
|
||||
{
|
||||
dtypes.push(ty.clone());
|
||||
} else {
|
||||
panic!("not support data format, {} - {}", file!(), line!());
|
||||
}
|
||||
}
|
||||
dtypes
|
||||
}
|
||||
|
||||
pub fn input_shapes(&self) -> &Vec<Vec<i64>> {
|
||||
&self.inputs.shapes
|
||||
}
|
||||
|
||||
pub fn input_names(&self) -> &Vec<String> {
|
||||
&self.inputs.names
|
||||
}
|
||||
|
||||
pub fn input_dtypes(&self) -> &Vec<TensorElementType> {
|
||||
&self.inputs.dtypes
|
||||
}
|
||||
|
||||
pub fn dtype(&self) -> TensorElementType {
|
||||
self.input_dtypes()[0]
|
||||
}
|
||||
|
||||
pub fn height(&self) -> u32 {
|
||||
self.inputs.sizes[0][0]
|
||||
}
|
||||
|
||||
pub fn width(&self) -> u32 {
|
||||
self.inputs.sizes[0][1]
|
||||
}
|
||||
|
||||
pub fn is_height_dynamic(&self) -> bool {
|
||||
self.input_shapes()[0][2] == -1
|
||||
}
|
||||
|
||||
pub fn is_width_dynamic(&self) -> bool {
|
||||
self.input_shapes()[0][3] == -1
|
||||
}
|
||||
|
||||
pub fn batch(&self) -> u32 {
|
||||
self.batch.opt
|
||||
}
|
||||
|
||||
pub fn is_batch_dynamic(&self) -> bool {
|
||||
self.input_shapes()[0][0] == -1
|
||||
}
|
||||
|
||||
pub fn ep(&self) -> &OrtEP {
|
||||
&self.ep
|
||||
}
|
||||
|
||||
pub fn task(&self) -> YOLOTask {
|
||||
self.task.clone()
|
||||
}
|
||||
|
||||
pub fn names(&self) -> Option<Vec<String>> {
|
||||
// class names, metadata parsing
|
||||
// String format: `{0: 'person', 1: 'bicycle', 2: 'sports ball', ..., 27: "yellow_lady's_slipper"}`
|
||||
match self.fetch_from_metadata("names") {
|
||||
Some(names) => {
|
||||
let re = Regex::new(r#"(['"])([-()\w '"]+)(['"])"#).unwrap();
|
||||
let mut names_ = vec![];
|
||||
for (_, [_, name, _]) in re.captures_iter(&names).map(|x| x.extract()) {
|
||||
names_.push(name.to_string());
|
||||
}
|
||||
Some(names_)
|
||||
}
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn nk(&self) -> Option<u32> {
|
||||
// num_keypoints, metadata parsing: String `nk` in onnx model: `[17, 3]`
|
||||
match self.fetch_from_metadata("kpt_shape") {
|
||||
None => None,
|
||||
Some(kpt_string) => {
|
||||
let re = Regex::new(r"([0-9]+), ([0-9]+)").unwrap();
|
||||
let caps = re.captures(&kpt_string).unwrap();
|
||||
Some(caps.get(1).unwrap().as_str().parse::<u32>().unwrap())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn nc(&self) -> Option<u32> {
|
||||
// num_classes
|
||||
match self.names() {
|
||||
// by names
|
||||
Some(names) => Some(names.len() as u32),
|
||||
None => match self.task() {
|
||||
// by task calculation
|
||||
YOLOTask::Classify => Some(self.output_shapes()[0][1] as u32),
|
||||
YOLOTask::Detect => {
|
||||
if self.output_shapes()[0][1] == -1 {
|
||||
None
|
||||
} else {
|
||||
// cxywhclss
|
||||
Some(self.output_shapes()[0][1] as u32 - 4)
|
||||
}
|
||||
}
|
||||
YOLOTask::Pose => {
|
||||
match self.nk() {
|
||||
None => None,
|
||||
Some(nk) => {
|
||||
if self.output_shapes()[0][1] == -1 {
|
||||
None
|
||||
} else {
|
||||
// cxywhclss3*kpt
|
||||
Some(self.output_shapes()[0][1] as u32 - 4 - 3 * nk)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
YOLOTask::Segment => {
|
||||
if self.output_shapes()[0][1] == -1 {
|
||||
None
|
||||
} else {
|
||||
// cxywhclssnm
|
||||
Some((self.output_shapes()[0][1] - self.output_shapes()[1][1]) as u32 - 4)
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn nm(&self) -> Option<u32> {
|
||||
// num_masks
|
||||
match self.task() {
|
||||
YOLOTask::Segment => Some(self.output_shapes()[1][1] as u32),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn na(&self) -> Option<u32> {
|
||||
// num_anchors
|
||||
match self.task() {
|
||||
YOLOTask::Segment | YOLOTask::Detect | YOLOTask::Pose => {
|
||||
if self.output_shapes()[0][2] == -1 {
|
||||
None
|
||||
} else {
|
||||
Some(self.output_shapes()[0][2] as u32)
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn author(&self) -> Option<String> {
|
||||
self.fetch_from_metadata("author")
|
||||
}
|
||||
|
||||
pub fn version(&self) -> Option<String> {
|
||||
self.fetch_from_metadata("version")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
use ndarray::{Array, Axis, IxDyn};
|
||||
|
||||
#[derive(Clone, PartialEq, Default)]
|
||||
pub struct YOLOResult {
|
||||
// YOLO tasks results of an image
|
||||
pub probs: Option<Embedding>,
|
||||
pub bboxes: Option<Vec<Bbox>>,
|
||||
pub keypoints: Option<Vec<Vec<Point2>>>,
|
||||
pub masks: Option<Vec<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for YOLOResult {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("YOLOResult")
|
||||
.field(
|
||||
"Probs(top5)",
|
||||
&format_args!("{:?}", self.probs().map(|probs| probs.topk(5))),
|
||||
)
|
||||
.field("Bboxes", &self.bboxes)
|
||||
.field("Keypoints", &self.keypoints)
|
||||
.field(
|
||||
"Masks",
|
||||
&format_args!("{:?}", self.masks().map(|masks| masks.len())),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl YOLOResult {
|
||||
pub fn new(
|
||||
probs: Option<Embedding>,
|
||||
bboxes: Option<Vec<Bbox>>,
|
||||
keypoints: Option<Vec<Vec<Point2>>>,
|
||||
masks: Option<Vec<Vec<u8>>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
probs,
|
||||
bboxes,
|
||||
keypoints,
|
||||
masks,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn probs(&self) -> Option<&Embedding> {
|
||||
self.probs.as_ref()
|
||||
}
|
||||
|
||||
pub fn keypoints(&self) -> Option<&Vec<Vec<Point2>>> {
|
||||
self.keypoints.as_ref()
|
||||
}
|
||||
|
||||
pub fn masks(&self) -> Option<&Vec<Vec<u8>>> {
|
||||
self.masks.as_ref()
|
||||
}
|
||||
|
||||
pub fn bboxes(&self) -> Option<&Vec<Bbox>> {
|
||||
self.bboxes.as_ref()
|
||||
}
|
||||
|
||||
pub fn bboxes_mut(&mut self) -> Option<&mut Vec<Bbox>> {
|
||||
self.bboxes.as_mut()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Clone, Default)]
|
||||
pub struct Point2 {
|
||||
// A point2d with x, y, conf
|
||||
x: f32,
|
||||
y: f32,
|
||||
confidence: f32,
|
||||
}
|
||||
|
||||
impl Point2 {
|
||||
pub fn new_with_conf(x: f32, y: f32, confidence: f32) -> Self {
|
||||
Self { x, y, confidence }
|
||||
}
|
||||
|
||||
pub fn new(x: f32, y: f32) -> Self {
|
||||
Self {
|
||||
x,
|
||||
y,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn x(&self) -> f32 {
|
||||
self.x
|
||||
}
|
||||
|
||||
pub fn y(&self) -> f32 {
|
||||
self.y
|
||||
}
|
||||
|
||||
pub fn confidence(&self) -> f32 {
|
||||
self.confidence
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default)]
|
||||
pub struct Embedding {
|
||||
// An float32 n-dims tensor
|
||||
data: Array<f32, IxDyn>,
|
||||
}
|
||||
|
||||
impl Embedding {
|
||||
pub fn new(data: Array<f32, IxDyn>) -> Self {
|
||||
Self { data }
|
||||
}
|
||||
|
||||
pub fn data(&self) -> &Array<f32, IxDyn> {
|
||||
&self.data
|
||||
}
|
||||
|
||||
pub fn topk(&self, k: usize) -> Vec<(usize, f32)> {
|
||||
let mut probs = self
|
||||
.data
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(a, b)| (a, *b))
|
||||
.collect::<Vec<_>>();
|
||||
probs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
let mut topk = Vec::new();
|
||||
for &(id, confidence) in probs.iter().take(k) {
|
||||
topk.push((id, confidence));
|
||||
}
|
||||
topk
|
||||
}
|
||||
|
||||
pub fn norm(&self) -> Array<f32, IxDyn> {
|
||||
let std_ = self.data.mapv(|x| x * x).sum_axis(Axis(0)).mapv(f32::sqrt);
|
||||
self.data.clone() / std_
|
||||
}
|
||||
|
||||
pub fn top1(&self) -> (usize, f32) {
|
||||
self.topk(1)[0]
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default)]
|
||||
pub struct Bbox {
|
||||
// a bounding box around an object
|
||||
xmin: f32,
|
||||
ymin: f32,
|
||||
width: f32,
|
||||
height: f32,
|
||||
id: usize,
|
||||
confidence: f32,
|
||||
}
|
||||
|
||||
impl Bbox {
|
||||
pub fn new_from_xywh(xmin: f32, ymin: f32, width: f32, height: f32) -> Self {
|
||||
Self {
|
||||
xmin,
|
||||
ymin,
|
||||
width,
|
||||
height,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new(xmin: f32, ymin: f32, width: f32, height: f32, id: usize, confidence: f32) -> Self {
|
||||
Self {
|
||||
xmin,
|
||||
ymin,
|
||||
width,
|
||||
height,
|
||||
id,
|
||||
confidence,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn width(&self) -> f32 {
|
||||
self.width
|
||||
}
|
||||
|
||||
pub fn height(&self) -> f32 {
|
||||
self.height
|
||||
}
|
||||
|
||||
pub fn xmin(&self) -> f32 {
|
||||
self.xmin
|
||||
}
|
||||
|
||||
pub fn ymin(&self) -> f32 {
|
||||
self.ymin
|
||||
}
|
||||
|
||||
pub fn xmax(&self) -> f32 {
|
||||
self.xmin + self.width
|
||||
}
|
||||
|
||||
pub fn ymax(&self) -> f32 {
|
||||
self.ymin + self.height
|
||||
}
|
||||
|
||||
pub fn tl(&self) -> Point2 {
|
||||
Point2::new(self.xmin, self.ymin)
|
||||
}
|
||||
|
||||
pub fn br(&self) -> Point2 {
|
||||
Point2::new(self.xmax(), self.ymax())
|
||||
}
|
||||
|
||||
pub fn cxcy(&self) -> Point2 {
|
||||
Point2::new(self.xmin + self.width / 2., self.ymin + self.height / 2.)
|
||||
}
|
||||
|
||||
pub fn id(&self) -> usize {
|
||||
self.id
|
||||
}
|
||||
|
||||
pub fn confidence(&self) -> f32 {
|
||||
self.confidence
|
||||
}
|
||||
|
||||
pub fn area(&self) -> f32 {
|
||||
self.width * self.height
|
||||
}
|
||||
|
||||
pub fn intersection_area(&self, another: &Bbox) -> f32 {
|
||||
let l = self.xmin.max(another.xmin);
|
||||
let r = (self.xmin + self.width).min(another.xmin + another.width);
|
||||
let t = self.ymin.max(another.ymin);
|
||||
let b = (self.ymin + self.height).min(another.ymin + another.height);
|
||||
(r - l + 1.).max(0.) * (b - t + 1.).max(0.)
|
||||
}
|
||||
|
||||
pub fn union(&self, another: &Bbox) -> f32 {
|
||||
self.area() + another.area() - self.intersection_area(another)
|
||||
}
|
||||
|
||||
pub fn iou(&self, another: &Bbox) -> f32 {
|
||||
let union = self.union(another);
|
||||
if union <= 0.0 {
|
||||
0.0
|
||||
} else {
|
||||
self.intersection_area(another) / union
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
# YOLOv8 - ONNX Runtime
|
||||
|
||||
This repository provides an example implementation for running [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) models using the [ONNX Runtime](https://onnxruntime.ai/). This allows for efficient inference across various hardware platforms supporting the [ONNX format](https://onnx.ai/).
|
||||
|
||||
## ⚙️ Installation
|
||||
|
||||
To get started, you'll need [Python](https://www.python.org/) installed. Then, install the necessary dependencies.
|
||||
|
||||
### Installing Required Dependencies
|
||||
|
||||
Clone the repository and install the packages listed in the `requirements.txt` file using [pip](https://pip.pypa.io/en/stable/):
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ultralytics/ultralytics.git
|
||||
cd ultralytics/examples/YOLOv8-ONNXRuntime
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### Installing ONNX Runtime Backend
|
||||
|
||||
You need to choose the appropriate ONNX Runtime package based on your hardware.
|
||||
|
||||
**GPU Acceleration (NVIDIA)**
|
||||
|
||||
If you have an NVIDIA GPU and want to leverage CUDA for faster inference, install the `onnxruntime-gpu` package. Ensure you have the correct [NVIDIA drivers](https://www.nvidia.com/Download/index.aspx) and CUDA toolkit installed. Refer to the official [ONNX Runtime GPU documentation](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html) for compatibility details.
|
||||
|
||||
```bash
|
||||
pip install onnxruntime-gpu
|
||||
```
|
||||
|
||||
**CPU Only**
|
||||
|
||||
If you don't have a compatible NVIDIA GPU or prefer CPU-based inference, install the standard `onnxruntime` package. Check the [ONNX Runtime installation guide](https://onnxruntime.ai/docs/install/) for more options.
|
||||
|
||||
```bash
|
||||
pip install onnxruntime
|
||||
```
|
||||
|
||||
## 🚀 Usage
|
||||
|
||||
Once the dependencies and the appropriate ONNX Runtime backend are installed, you can perform inference using the provided Python script.
|
||||
|
||||
### Exporting Your Model
|
||||
|
||||
Before running inference, you need a YOLOv8 model in the ONNX format (`.onnx`). You can export your trained Ultralytics YOLOv8 models using the Ultralytics CLI or Python SDK. See the [Ultralytics export documentation](https://docs.ultralytics.com/modes/export/) for detailed instructions.
|
||||
|
||||
Example export command:
|
||||
|
||||
```bash
|
||||
yolo export model=yolov8n.pt format=onnx # Export yolov8n model to ONNX
|
||||
```
|
||||
|
||||
### Running Inference
|
||||
|
||||
Execute the `main.py` script with the path to your ONNX model and input image. You can also adjust the confidence and [Intersection over Union (IoU)](https://www.ultralytics.com/glossary/intersection-over-union-iou) thresholds for [object detection](https://docs.ultralytics.com/tasks/detect/).
|
||||
|
||||
```bash
|
||||
python main.py --model yolov8n.onnx --img image.jpg --conf-thres 0.5 --iou-thres 0.5
|
||||
```
|
||||
|
||||
- `--model`: Path to the YOLOv8 ONNX model file (e.g., `yolov8n.onnx`).
|
||||
- `--img`: Path to the input image (e.g., `image.jpg`).
|
||||
- `--conf-thres`: Confidence threshold for filtering detections. Only detections with a score higher than this value will be kept. Learn more about thresholds in the [performance metrics guide](https://docs.ultralytics.com/guides/yolo-performance-metrics/).
|
||||
- `--iou-thres`: IoU threshold for Non-Maximum Suppression (NMS). Boxes with IoU greater than this threshold will be suppressed. See the [NMS glossary entry](https://www.ultralytics.com/glossary/non-maximum-suppression-nms) for details.
|
||||
|
||||
The script will process the image, perform object detection, draw bounding boxes on the detected objects, and save the output image as `output.jpg`.
|
||||
|
||||
## Contributing
|
||||
|
||||
Contributions to enhance this example or add new features are welcome! Please refer to the main [Ultralytics repository](https://github.com/ultralytics/ultralytics) for contribution guidelines. If you encounter issues or have suggestions, feel free to open an issue on the [ONNX Runtime GitHub](https://github.com/microsoft/onnxruntime) or the Ultralytics repository.
|
||||
282
algorithms/dms_yolo/code/examples/YOLOv8-ONNXRuntime/main.py
Normal file
282
algorithms/dms_yolo/code/examples/YOLOv8-ONNXRuntime/main.py
Normal file
@@ -0,0 +1,282 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import torch
|
||||
|
||||
from ultralytics.utils import ASSETS, YAML
|
||||
from ultralytics.utils.checks import check_requirements, check_yaml
|
||||
|
||||
|
||||
class YOLOv8:
|
||||
"""YOLOv8 object detection model class for handling ONNX inference and visualization.
|
||||
|
||||
This class provides functionality to load a YOLOv8 ONNX model, perform inference on images, and visualize the
|
||||
detection results with bounding boxes and labels.
|
||||
|
||||
Attributes:
|
||||
onnx_model (str): Path to the ONNX model file.
|
||||
input_image (str): Path to the input image file.
|
||||
confidence_thres (float): Confidence threshold for filtering detections.
|
||||
iou_thres (float): IoU threshold for non-maximum suppression.
|
||||
classes (list[str]): List of class names from the COCO dataset.
|
||||
color_palette (np.ndarray): Random color palette for visualizing different classes.
|
||||
input_width (int): Width dimension of the model input.
|
||||
input_height (int): Height dimension of the model input.
|
||||
img (np.ndarray): The loaded input image.
|
||||
img_height (int): Height of the input image.
|
||||
img_width (int): Width of the input image.
|
||||
|
||||
Methods:
|
||||
letterbox: Resize and reshape images while maintaining aspect ratio by adding padding.
|
||||
draw_detections: Draw bounding boxes and labels on the input image based on detected objects.
|
||||
preprocess: Preprocess the input image before performing inference.
|
||||
postprocess: Perform post-processing on the model's output to extract and visualize detections.
|
||||
main: Perform inference using an ONNX model and return the output image with drawn detections.
|
||||
|
||||
Examples:
|
||||
Initialize YOLOv8 detector and run inference
|
||||
>>> detector = YOLOv8("yolov8n.onnx", "image.jpg", 0.5, 0.5)
|
||||
>>> output_image = detector.main()
|
||||
"""
|
||||
|
||||
def __init__(self, onnx_model: str, input_image: str, confidence_thres: float, iou_thres: float):
|
||||
"""Initialize an instance of the YOLOv8 class.
|
||||
|
||||
Args:
|
||||
onnx_model (str): Path to the ONNX model.
|
||||
input_image (str): Path to the input image.
|
||||
confidence_thres (float): Confidence threshold for filtering detections.
|
||||
iou_thres (float): IoU threshold for non-maximum suppression.
|
||||
"""
|
||||
self.onnx_model = onnx_model
|
||||
self.input_image = input_image
|
||||
self.confidence_thres = confidence_thres
|
||||
self.iou_thres = iou_thres
|
||||
|
||||
# Load the class names from the COCO dataset
|
||||
self.classes = YAML.load(check_yaml("coco8.yaml"))["names"]
|
||||
|
||||
# Generate a color palette for the classes
|
||||
self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))
|
||||
|
||||
def letterbox(self, img: np.ndarray, new_shape: tuple[int, int] = (640, 640)) -> tuple[np.ndarray, tuple[int, int]]:
|
||||
"""Resize and reshape images while maintaining aspect ratio by adding padding.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): Input image to be resized.
|
||||
new_shape (tuple[int, int]): Target shape (height, width) for the image.
|
||||
|
||||
Returns:
|
||||
img (np.ndarray): Resized and padded image.
|
||||
pad (tuple[int, int]): Padding values (top, left) applied to the image.
|
||||
"""
|
||||
shape = img.shape[:2] # current shape [height, width]
|
||||
|
||||
# Scale ratio (new / old)
|
||||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||
|
||||
# Compute padding
|
||||
new_unpad = round(shape[1] * r), round(shape[0] * r)
|
||||
dw, dh = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2 # wh padding
|
||||
|
||||
if shape[::-1] != new_unpad: # resize
|
||||
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||
top, bottom = round(dh - 0.1), round(dh + 0.1)
|
||||
left, right = round(dw - 0.1), round(dw + 0.1)
|
||||
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
|
||||
|
||||
return img, (top, left)
|
||||
|
||||
def draw_detections(self, img: np.ndarray, box: list[float], score: float, class_id: int) -> None:
|
||||
"""Draw bounding boxes and labels on the input image based on the detected objects."""
|
||||
# Extract the coordinates of the bounding box
|
||||
x1, y1, w, h = box
|
||||
|
||||
# Retrieve the color for the class ID
|
||||
color = self.color_palette[class_id]
|
||||
|
||||
# Draw the bounding box on the image
|
||||
cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)
|
||||
|
||||
# Create the label text with class name and score
|
||||
label = f"{self.classes[class_id]}: {score:.2f}"
|
||||
|
||||
# Calculate the dimensions of the label text
|
||||
(label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||||
|
||||
# Calculate the position of the label text
|
||||
label_x = x1
|
||||
label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10
|
||||
|
||||
# Draw a filled rectangle as the background for the label text
|
||||
cv2.rectangle(
|
||||
img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
|
||||
)
|
||||
|
||||
# Draw the label text on the image
|
||||
cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
|
||||
|
||||
def preprocess(self) -> tuple[np.ndarray, tuple[int, int]]:
|
||||
"""Preprocess the input image before performing inference.
|
||||
|
||||
This method reads the input image, converts its color space, applies letterboxing to maintain aspect ratio,
|
||||
normalizes pixel values, and prepares the image data for model input.
|
||||
|
||||
Returns:
|
||||
image_data (np.ndarray): Preprocessed image data ready for inference with shape (1, 3, height, width).
|
||||
pad (tuple[int, int]): Padding values (top, left) applied during letterboxing.
|
||||
"""
|
||||
# Read the input image using OpenCV
|
||||
self.img = cv2.imread(self.input_image)
|
||||
|
||||
# Get the height and width of the input image
|
||||
self.img_height, self.img_width = self.img.shape[:2]
|
||||
|
||||
# Convert the image color space from BGR to RGB
|
||||
img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
img, pad = self.letterbox(img, (self.input_width, self.input_height))
|
||||
|
||||
# Normalize the image data by dividing it by 255.0
|
||||
image_data = np.array(img) / 255.0
|
||||
|
||||
# Transpose the image to have the channel dimension as the first dimension
|
||||
image_data = np.transpose(image_data, (2, 0, 1)) # Channel first
|
||||
|
||||
# Expand the dimensions of the image data to match the expected input shape
|
||||
image_data = image_data[None].astype(np.float32)
|
||||
|
||||
# Return the preprocessed image data
|
||||
return image_data, pad
|
||||
|
||||
def postprocess(self, input_image: np.ndarray, output: list[np.ndarray], pad: tuple[int, int]) -> np.ndarray:
|
||||
"""Perform post-processing on the model's output to extract and visualize detections.
|
||||
|
||||
This method processes the raw model output to extract bounding boxes, scores, and class IDs. It applies
|
||||
non-maximum suppression to filter overlapping detections and draws the results on the input image.
|
||||
|
||||
Args:
|
||||
input_image (np.ndarray): The input image.
|
||||
output (list[np.ndarray]): The output arrays from the model.
|
||||
pad (tuple[int, int]): Padding values (top, left) used during letterboxing.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The input image with detections drawn on it.
|
||||
"""
|
||||
# Transpose and squeeze the output to match the expected shape
|
||||
outputs = np.transpose(np.squeeze(output[0]))
|
||||
|
||||
# Get the number of rows in the outputs array
|
||||
rows = outputs.shape[0]
|
||||
|
||||
# Lists to store the bounding boxes, scores, and class IDs of the detections
|
||||
boxes = []
|
||||
scores = []
|
||||
class_ids = []
|
||||
|
||||
# Calculate the scaling factors for the bounding box coordinates
|
||||
gain = min(self.input_height / self.img_height, self.input_width / self.img_width)
|
||||
outputs[:, 0] -= pad[1]
|
||||
outputs[:, 1] -= pad[0]
|
||||
|
||||
# Iterate over each row in the outputs array
|
||||
for i in range(rows):
|
||||
# Extract the class scores from the current row
|
||||
classes_scores = outputs[i][4:]
|
||||
|
||||
# Find the maximum score among the class scores
|
||||
max_score = np.amax(classes_scores)
|
||||
|
||||
# If the maximum score is above the confidence threshold
|
||||
if max_score >= self.confidence_thres:
|
||||
# Get the class ID with the highest score
|
||||
class_id = np.argmax(classes_scores)
|
||||
|
||||
# Extract the bounding box coordinates from the current row
|
||||
x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
|
||||
|
||||
# Calculate the scaled coordinates of the bounding box
|
||||
left = int((x - w / 2) / gain)
|
||||
top = int((y - h / 2) / gain)
|
||||
width = int(w / gain)
|
||||
height = int(h / gain)
|
||||
|
||||
# Add the class ID, score, and box coordinates to the respective lists
|
||||
class_ids.append(class_id)
|
||||
scores.append(max_score)
|
||||
boxes.append([left, top, width, height])
|
||||
|
||||
# Apply non-maximum suppression to filter out overlapping bounding boxes
|
||||
indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)
|
||||
|
||||
# Iterate over the selected indices after non-maximum suppression
|
||||
for i in np.array(indices).flatten():
|
||||
# Get the box, score, and class ID corresponding to the index
|
||||
box = boxes[int(i)]
|
||||
score = scores[int(i)]
|
||||
class_id = class_ids[int(i)]
|
||||
|
||||
# Draw the detection on the input image
|
||||
self.draw_detections(input_image, box, score, class_id)
|
||||
|
||||
# Return the modified input image
|
||||
return input_image
|
||||
|
||||
def main(self) -> np.ndarray:
|
||||
"""Perform inference using an ONNX model and return the output image with drawn detections.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The output image with drawn detections.
|
||||
"""
|
||||
available = ort.get_available_providers()
|
||||
providers = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in available]
|
||||
session = ort.InferenceSession(self.onnx_model, providers=providers or available)
|
||||
|
||||
# Get the model inputs
|
||||
model_inputs = session.get_inputs()
|
||||
|
||||
# Store the shape of the input for later use
|
||||
input_shape = model_inputs[0].shape
|
||||
self.input_height = input_shape[2]
|
||||
self.input_width = input_shape[3]
|
||||
|
||||
# Preprocess the image data
|
||||
img_data, pad = self.preprocess()
|
||||
|
||||
# Run inference using the preprocessed image data
|
||||
outputs = session.run(None, {model_inputs[0].name: img_data})
|
||||
|
||||
# Perform post-processing on the outputs to obtain output image
|
||||
return self.postprocess(self.img, outputs, pad)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create an argument parser to handle command-line arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", type=str, default="yolov8n.onnx", help="Input your ONNX model.")
|
||||
parser.add_argument("--img", type=str, default=str(ASSETS / "bus.jpg"), help="Path to input image.")
|
||||
parser.add_argument("--conf-thres", type=float, default=0.5, help="Confidence threshold")
|
||||
parser.add_argument("--iou-thres", type=float, default=0.5, help="NMS IoU threshold")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Check the requirements and select the appropriate backend (CPU or GPU)
|
||||
check_requirements("onnxruntime-gpu" if torch.cuda.is_available() else "onnxruntime")
|
||||
|
||||
# Create an instance of the YOLOv8 class with the specified arguments
|
||||
detection = YOLOv8(args.model, args.img, args.conf_thres, args.iou_thres)
|
||||
|
||||
# Perform object detection and obtain the output image
|
||||
output_image = detection.main()
|
||||
|
||||
# Display the output image in a window
|
||||
cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
|
||||
cv2.imshow("Output", output_image)
|
||||
|
||||
# Wait for a key press to exit
|
||||
cv2.waitKey(0)
|
||||
@@ -0,0 +1,4 @@
|
||||
numpy
|
||||
opencv-python
|
||||
onnxruntime
|
||||
ultralytics
|
||||
@@ -0,0 +1,51 @@
|
||||
# Ultralytics YOLOv8 Object Detection with OpenCV and ONNX
|
||||
|
||||
This example demonstrates how to implement [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) object detection using [OpenCV](https://opencv.org/) in [Python](https://www.python.org/), leveraging the [ONNX (Open Neural Network Exchange)](https://onnx.ai/) model format for efficient inference.
|
||||
|
||||
## 🚀 Getting Started
|
||||
|
||||
Follow these simple steps to get the example running on your local machine.
|
||||
|
||||
1. **Clone the Repository:**
|
||||
If you haven't already, clone the Ultralytics repository to access the example code:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ultralytics/ultralytics.git
|
||||
cd ultralytics/examples/YOLOv8-OpenCV-ONNX-Python/
|
||||
```
|
||||
|
||||
2. **Install Requirements:**
|
||||
Install the necessary Python packages listed in the `requirements.txt` file. We recommend using a virtual environment.
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
3. **Run the Detection Script:**
|
||||
Execute the main Python script, specifying the ONNX model path and the input image.
|
||||
```bash
|
||||
python main.py --model yolov8n.onnx --img image.jpg
|
||||
```
|
||||
The script will perform object detection on `image.jpg` using the `yolov8n.onnx` model and display the results.
|
||||
|
||||
## 🛠️ Exporting Your Model
|
||||
|
||||
If you want to use a different Ultralytics YOLOv8 model or one you've trained yourself, you need to export it to the ONNX format first.
|
||||
|
||||
1. **Install Ultralytics:**
|
||||
If you don't have it installed, get the latest `ultralytics` package:
|
||||
|
||||
```bash
|
||||
pip install ultralytics
|
||||
```
|
||||
|
||||
2. **Export the Model:**
|
||||
Use the `yolo export` command to convert your desired model (e.g., `yolov8n.pt`) to ONNX. Ensure you specify `opset=12` or higher for compatibility with OpenCV's DNN module. You can find more details in the Ultralytics [Export documentation](https://docs.ultralytics.com/modes/export/).
|
||||
```bash
|
||||
yolo export model=yolov8n.pt imgsz=640 format=onnx opset=12
|
||||
```
|
||||
This command will generate a `yolov8n.onnx` file (or the corresponding name for your model) in your working directory. You can then use this `.onnx` file with the `main.py` script.
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! If you find any issues or have suggestions for improvement, please feel free to open an issue or submit a pull request to the main [Ultralytics repository](https://github.com/ultralytics/ultralytics). Thank you for helping us make Ultralytics YOLO even better!
|
||||
@@ -0,0 +1,134 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from typing import Any
|
||||
|
||||
import cv2.dnn
|
||||
import numpy as np
|
||||
|
||||
from ultralytics.utils import ASSETS, YAML
|
||||
from ultralytics.utils.checks import check_yaml
|
||||
|
||||
CLASSES = YAML.load(check_yaml("coco8.yaml"))["names"]
|
||||
colors = np.random.uniform(0, 255, size=(len(CLASSES), 3))
|
||||
|
||||
|
||||
def draw_bounding_box(
|
||||
img: np.ndarray, class_id: int, confidence: float, x: int, y: int, x_plus_w: int, y_plus_h: int
|
||||
) -> None:
|
||||
"""Draw bounding boxes on the input image based on the provided arguments.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): The input image to draw the bounding box on.
|
||||
class_id (int): Class ID of the detected object.
|
||||
confidence (float): Confidence score of the detected object.
|
||||
x (int): X-coordinate of the top-left corner of the bounding box.
|
||||
y (int): Y-coordinate of the top-left corner of the bounding box.
|
||||
x_plus_w (int): X-coordinate of the bottom-right corner of the bounding box.
|
||||
y_plus_h (int): Y-coordinate of the bottom-right corner of the bounding box.
|
||||
"""
|
||||
label = f"{CLASSES[class_id]} ({confidence:.2f})"
|
||||
color = colors[class_id]
|
||||
cv2.rectangle(img, (x, y), (x_plus_w, y_plus_h), color, 2)
|
||||
cv2.putText(img, label, (x - 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
||||
|
||||
|
||||
def main(onnx_model: str, input_image: str) -> list[dict[str, Any]]:
|
||||
"""Load ONNX model, perform inference, draw bounding boxes, and display the output image.
|
||||
|
||||
Args:
|
||||
onnx_model (str): Path to the ONNX model.
|
||||
input_image (str): Path to the input image.
|
||||
|
||||
Returns:
|
||||
(list[dict[str, Any]]): List of dictionaries containing detection information such as class_id, class_name,
|
||||
confidence, box coordinates, and scale factor.
|
||||
"""
|
||||
# Load the ONNX model
|
||||
model: cv2.dnn.Net = cv2.dnn.readNetFromONNX(onnx_model)
|
||||
|
||||
# Read the input image
|
||||
original_image: np.ndarray = cv2.imread(input_image)
|
||||
[height, width, _] = original_image.shape
|
||||
|
||||
# Prepare a square image for inference
|
||||
length = max((height, width))
|
||||
image = np.zeros((length, length, 3), np.uint8)
|
||||
image[0:height, 0:width] = original_image
|
||||
|
||||
# Calculate scale factor
|
||||
scale = length / 640
|
||||
|
||||
# Preprocess the image and prepare blob for model
|
||||
blob = cv2.dnn.blobFromImage(image, scalefactor=1 / 255, size=(640, 640), swapRB=True)
|
||||
model.setInput(blob)
|
||||
|
||||
# Perform inference
|
||||
outputs = model.forward()
|
||||
|
||||
# Prepare output array
|
||||
outputs = np.array([cv2.transpose(outputs[0])])
|
||||
rows = outputs.shape[1]
|
||||
|
||||
boxes = []
|
||||
scores = []
|
||||
class_ids = []
|
||||
|
||||
# Iterate through output to collect bounding boxes, confidence scores, and class IDs
|
||||
for i in range(rows):
|
||||
classes_scores = outputs[0][i][4:]
|
||||
(_minScore, maxScore, _minClassLoc, (_x, maxClassIndex)) = cv2.minMaxLoc(classes_scores)
|
||||
if maxScore >= 0.25:
|
||||
box = [
|
||||
outputs[0][i][0] - (0.5 * outputs[0][i][2]), # x center - width/2 = left x
|
||||
outputs[0][i][1] - (0.5 * outputs[0][i][3]), # y center - height/2 = top y
|
||||
outputs[0][i][2], # width
|
||||
outputs[0][i][3], # height
|
||||
]
|
||||
boxes.append(box)
|
||||
scores.append(maxScore)
|
||||
class_ids.append(maxClassIndex)
|
||||
|
||||
# Apply NMS (Non-maximum suppression)
|
||||
result_boxes = np.array(cv2.dnn.NMSBoxes(boxes, scores, 0.25, 0.45, 0.5)).flatten()
|
||||
|
||||
detections = []
|
||||
|
||||
# Iterate through NMS results to draw bounding boxes and labels
|
||||
for index in result_boxes:
|
||||
index = int(index)
|
||||
box = boxes[index]
|
||||
detection = {
|
||||
"class_id": class_ids[index],
|
||||
"class_name": CLASSES[class_ids[index]],
|
||||
"confidence": scores[index],
|
||||
"box": box,
|
||||
"scale": scale,
|
||||
}
|
||||
detections.append(detection)
|
||||
draw_bounding_box(
|
||||
original_image,
|
||||
class_ids[index],
|
||||
scores[index],
|
||||
round(box[0] * scale),
|
||||
round(box[1] * scale),
|
||||
round((box[0] + box[2]) * scale),
|
||||
round((box[1] + box[3]) * scale),
|
||||
)
|
||||
|
||||
# Display the image with bounding boxes
|
||||
cv2.imshow("image", original_image)
|
||||
cv2.waitKey(0)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
return detections
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", default="yolov8n.onnx", help="Input your ONNX model.")
|
||||
parser.add_argument("--img", default=str(ASSETS / "bus.jpg"), help="Path to input image.")
|
||||
args = parser.parse_args()
|
||||
main(args.model, args.img)
|
||||
@@ -0,0 +1,3 @@
|
||||
numpy
|
||||
opencv-python
|
||||
ultralytics
|
||||
@@ -0,0 +1,21 @@
|
||||
cmake_minimum_required(VERSION 3.12)
|
||||
project(yolov8_openvino_example)
|
||||
|
||||
set(CMAKE_CXX_STANDARD 14)
|
||||
|
||||
find_package(OpenCV REQUIRED)
|
||||
|
||||
include_directories(
|
||||
${OpenCV_INCLUDE_DIRS}
|
||||
/path/to/intel/openvino/runtime/include
|
||||
)
|
||||
|
||||
add_executable(detect
|
||||
main.cc
|
||||
inference.cc
|
||||
)
|
||||
|
||||
target_link_libraries(detect
|
||||
${OpenCV_LIBS}
|
||||
/path/to/intel/openvino/runtime/lib/intel64/libopenvino.so
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
# YOLOv8 OpenVINO Inference in C++
|
||||
|
||||
Welcome to the [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) OpenVINO Inference example in C++! This guide will help you get started with leveraging the powerful YOLOv8 models using the [Intel OpenVINO™ toolkit](https://docs.openvino.ai/) and [OpenCV API](https://docs.opencv.org/) in your C++ projects. Whether you're looking to enhance performance on Intel hardware or add flexibility to your applications, this example provides a solid foundation. Learn more about optimizing models on the [Ultralytics blog](https://www.ultralytics.com/blog).
|
||||
|
||||
## 🌟 Features
|
||||
|
||||
- 🚀 **Model Format Support**: Compatible with [ONNX](https://onnx.ai/) and [OpenVINO Intermediate Representation (IR)](https://docs.openvino.ai/2023.3/openvino_docs_MO_DG_IR_and_opsets.html) formats. Check the [Ultralytics ONNX integration](https://docs.ultralytics.com/integrations/onnx/) for more details.
|
||||
- ⚡ **Precision Options**: Run models in **FP32**, **FP16** ([half-precision](https://www.ultralytics.com/glossary/half-precision)), and **INT8** ([quantization](https://www.ultralytics.com/glossary/model-quantization)) precisions for optimized performance.
|
||||
- 🔄 **Dynamic Shape Loading**: Easily handle models with dynamic input shapes, common in many [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) tasks.
|
||||
|
||||
## 📋 Dependencies
|
||||
|
||||
To ensure smooth execution, please make sure you have the following dependencies installed:
|
||||
|
||||
| Dependency | Version |
|
||||
| ----------------------------------------------------- | -------- |
|
||||
| [OpenVINO](https://docs.openvino.ai/latest/home.html) | >=2023.3 |
|
||||
| [OpenCV](https://opencv.org/) | >=4.5.0 |
|
||||
| [C++](https://en.cppreference.com/w/) | >=14 |
|
||||
| [CMake](https://cmake.org/documentation/) | >=3.12.0 |
|
||||
|
||||
## ⚙️ Build Instructions
|
||||
|
||||
Follow these steps to build the project:
|
||||
|
||||
1. Clone the Ultralytics repository:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ultralytics/ultralytics.git
|
||||
cd ultralytics/examples/YOLOv8-OpenVINO-CPP-Inference
|
||||
```
|
||||
|
||||
2. Create a build directory and compile the project using CMake:
|
||||
```bash
|
||||
mkdir build
|
||||
cd build
|
||||
cmake ..
|
||||
make
|
||||
```
|
||||
|
||||
## 🛠️ Usage
|
||||
|
||||
Once built, you can run [inference](https://www.ultralytics.com/glossary/real-time-inference) on an image using the compiled executable. Provide the path to your model file (either `.xml` for OpenVINO IR or `.onnx`) and the path to your image:
|
||||
|
||||
```bash
|
||||
# Example using an OpenVINO IR model
|
||||
./detect path/to/your/model.xml path/to/your/image.jpg
|
||||
|
||||
# Example using an ONNX model
|
||||
./detect path/to/your/model.onnx path/to/your/image.jpg
|
||||
```
|
||||
|
||||
This command will process the image using the specified YOLOv8 model and display the [object detection](https://www.ultralytics.com/glossary/object-detection) results. Explore various [Ultralytics Solutions](https://docs.ultralytics.com/solutions/) for real-world applications.
|
||||
|
||||
## 🔄 Exporting YOLOv8 Models
|
||||
|
||||
To use your Ultralytics YOLOv8 model with this C++ example, you first need to export it to the OpenVINO IR or ONNX format. Use the `yolo export` command available in the Ultralytics Python package. Find detailed instructions in the [Export mode documentation](https://docs.ultralytics.com/modes/export/).
|
||||
|
||||
```bash
|
||||
# Export to OpenVINO format (generates .xml and .bin files)
|
||||
yolo export model=yolov8s.pt imgsz=640 format=openvino
|
||||
|
||||
# Export to ONNX format
|
||||
yolo export model=yolov8s.pt imgsz=640 format=onnx
|
||||
```
|
||||
|
||||
For more details on exporting and optimizing models for OpenVINO, refer to the [Ultralytics OpenVINO integration guide](https://docs.ultralytics.com/integrations/openvino/).
|
||||
|
||||
## 📸 Screenshots
|
||||
|
||||
### Running Using OpenVINO Model
|
||||
|
||||
|
||||
|
||||
### Running Using ONNX Model
|
||||
|
||||
|
||||
|
||||
## ❤️ Contributions
|
||||
|
||||
We hope this example helps you integrate YOLOv8 with OpenVINO and OpenCV into your C++ projects effortlessly. Contributions to improve this example or add new features are welcome! Please see the [Ultralytics contribution guidelines](https://docs.ultralytics.com/help/contributing/) for more information. Visit the main [Ultralytics documentation](https://docs.ultralytics.com/) for further guides and resources. Happy coding! 🚀
|
||||
@@ -0,0 +1,178 @@
|
||||
#include "inference.h"
|
||||
|
||||
#include <memory>
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <random>
|
||||
|
||||
namespace yolo {
|
||||
|
||||
// Constructor to initialize the model with default input shape
|
||||
Inference::Inference(const std::string &model_path, const float &model_confidence_threshold, const float &model_NMS_threshold) {
|
||||
model_input_shape_ = cv::Size(640, 640); // Set the default size for models with dynamic shapes to prevent errors.
|
||||
model_confidence_threshold_ = model_confidence_threshold;
|
||||
model_NMS_threshold_ = model_NMS_threshold;
|
||||
InitializeModel(model_path);
|
||||
}
|
||||
|
||||
// Constructor to initialize the model with specified input shape
|
||||
Inference::Inference(const std::string &model_path, const cv::Size model_input_shape, const float &model_confidence_threshold, const float &model_NMS_threshold) {
|
||||
model_input_shape_ = model_input_shape;
|
||||
model_confidence_threshold_ = model_confidence_threshold;
|
||||
model_NMS_threshold_ = model_NMS_threshold;
|
||||
InitializeModel(model_path);
|
||||
}
|
||||
|
||||
void Inference::InitializeModel(const std::string &model_path) {
|
||||
ov::Core core; // OpenVINO core object
|
||||
std::shared_ptr<ov::Model> model = core.read_model(model_path); // Read the model from file
|
||||
|
||||
// If the model has dynamic shapes, reshape it to the specified input shape
|
||||
if (model->is_dynamic()) {
|
||||
model->reshape({1, 3, static_cast<long int>(model_input_shape_.height), static_cast<long int>(model_input_shape_.width)});
|
||||
}
|
||||
|
||||
// Preprocessing setup for the model
|
||||
ov::preprocess::PrePostProcessor ppp = ov::preprocess::PrePostProcessor(model);
|
||||
ppp.input().tensor().set_element_type(ov::element::u8).set_layout("NHWC").set_color_format(ov::preprocess::ColorFormat::BGR);
|
||||
ppp.input().preprocess().convert_element_type(ov::element::f32).convert_color(ov::preprocess::ColorFormat::RGB).scale({255, 255, 255});
|
||||
ppp.input().model().set_layout("NCHW");
|
||||
ppp.output().tensor().set_element_type(ov::element::f32);
|
||||
model = ppp.build(); // Build the preprocessed model
|
||||
|
||||
// Compile the model for inference
|
||||
compiled_model_ = core.compile_model(model, "AUTO");
|
||||
inference_request_ = compiled_model_.create_infer_request(); // Create inference request
|
||||
|
||||
short width, height;
|
||||
|
||||
// Get input shape from the model
|
||||
const std::vector<ov::Output<ov::Node>> inputs = model->inputs();
|
||||
const ov::Shape input_shape = inputs[0].get_shape();
|
||||
height = input_shape[1];
|
||||
width = input_shape[2];
|
||||
model_input_shape_ = cv::Size2f(width, height);
|
||||
|
||||
// Get output shape from the model
|
||||
const std::vector<ov::Output<ov::Node>> outputs = model->outputs();
|
||||
const ov::Shape output_shape = outputs[0].get_shape();
|
||||
height = output_shape[1];
|
||||
width = output_shape[2];
|
||||
model_output_shape_ = cv::Size(width, height);
|
||||
}
|
||||
|
||||
// Method to run inference on an input frame
|
||||
void Inference::RunInference(cv::Mat &frame) {
|
||||
Preprocessing(frame); // Preprocess the input frame
|
||||
inference_request_.infer(); // Run inference
|
||||
PostProcessing(frame); // Postprocess the inference results
|
||||
}
|
||||
|
||||
// Method to preprocess the input frame
|
||||
void Inference::Preprocessing(const cv::Mat &frame) {
|
||||
cv::Mat resized_frame;
|
||||
cv::resize(frame, resized_frame, model_input_shape_, 0, 0, cv::INTER_AREA); // Resize the frame to match the model input shape
|
||||
|
||||
// Calculate scaling factor
|
||||
scale_factor_.x = static_cast<float>(frame.cols / model_input_shape_.width);
|
||||
scale_factor_.y = static_cast<float>(frame.rows / model_input_shape_.height);
|
||||
|
||||
ov::Tensor input_tensor = inference_request_.get_input_tensor();
|
||||
uint8_t* input_data = input_tensor.data<uint8_t>();
|
||||
size_t bytes_to_copy = resized_frame.total() * resized_frame.elemSize();
|
||||
memcpy(input_data, resized_frame.data, bytes_to_copy);
|
||||
|
||||
inference_request_.set_input_tensor(input_tensor); // Set input tensor for inference
|
||||
}
|
||||
|
||||
// Method to postprocess the inference results
|
||||
void Inference::PostProcessing(cv::Mat &frame) {
|
||||
std::vector<int> class_list;
|
||||
std::vector<float> confidence_list;
|
||||
std::vector<cv::Rect> box_list;
|
||||
|
||||
// Get the output tensor from the inference request
|
||||
const float *detections = inference_request_.get_output_tensor().data<const float>();
|
||||
const cv::Mat detection_outputs(model_output_shape_, CV_32F, (float *)detections); // Create OpenCV matrix from output tensor
|
||||
|
||||
// Iterate over detections and collect class IDs, confidence scores, and bounding boxes
|
||||
for (int i = 0; i < detection_outputs.cols; ++i) {
|
||||
const cv::Mat classes_scores = detection_outputs.col(i).rowRange(4, detection_outputs.rows);
|
||||
|
||||
cv::Point class_id;
|
||||
double score;
|
||||
cv::minMaxLoc(classes_scores, nullptr, &score, nullptr, &class_id); // Find the class with the highest score
|
||||
|
||||
// Check if the detection meets the confidence threshold
|
||||
if (score > model_confidence_threshold_) {
|
||||
class_list.push_back(class_id.y);
|
||||
confidence_list.push_back(score);
|
||||
|
||||
const float cx = detection_outputs.at<float>(0, i);
|
||||
const float cy = detection_outputs.at<float>(1, i);
|
||||
const float w = detection_outputs.at<float>(2, i);
|
||||
const float h = detection_outputs.at<float>(3, i);
|
||||
|
||||
cv::Rect box;
|
||||
box.x = static_cast<int>((cx - w / 2));
|
||||
box.y = static_cast<int>((cy - h / 2));
|
||||
box.width = static_cast<int>(w);
|
||||
box.height = static_cast<int>(h);
|
||||
box_list.push_back(box);
|
||||
}
|
||||
}
|
||||
|
||||
// Apply Non-Maximum Suppression (NMS) to filter overlapping bounding boxes
|
||||
std::vector<int> NMS_result;
|
||||
cv::dnn::NMSBoxes(box_list, confidence_list, model_confidence_threshold_, model_NMS_threshold_, NMS_result);
|
||||
|
||||
// Collect final detections after NMS
|
||||
for (int i = 0; i < NMS_result.size(); ++i) {
|
||||
Detection result;
|
||||
const unsigned short id = NMS_result[i];
|
||||
|
||||
result.class_id = class_list[id];
|
||||
result.confidence = confidence_list[id];
|
||||
result.box = GetBoundingBox(box_list[id]);
|
||||
|
||||
DrawDetectedObject(frame, result);
|
||||
}
|
||||
}
|
||||
|
||||
// Method to get the bounding box in the correct scale
|
||||
cv::Rect Inference::GetBoundingBox(const cv::Rect &src) const {
|
||||
cv::Rect box = src;
|
||||
box.x = static_cast<int>(box.x * scale_factor_.x);
|
||||
box.y = static_cast<int>(box.y * scale_factor_.y);
|
||||
box.width = static_cast<int>(box.width * scale_factor_.x);
|
||||
box.height = static_cast<int>(box.height * scale_factor_.y);
|
||||
return box;
|
||||
}
|
||||
|
||||
void Inference::DrawDetectedObject(cv::Mat &frame, const Detection &detection) const {
|
||||
const cv::Rect &box = detection.box;
|
||||
const float &confidence = detection.confidence;
|
||||
const int &class_id = detection.class_id;
|
||||
|
||||
// Generate a random color for the bounding box
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_int_distribution<int> dis(120, 255);
|
||||
const cv::Scalar &color = cv::Scalar(dis(gen), dis(gen), dis(gen));
|
||||
|
||||
// Draw the bounding box around the detected object
|
||||
cv::rectangle(frame, cv::Point(box.x, box.y), cv::Point(box.x + box.width, box.y + box.height), color, 3);
|
||||
|
||||
// Prepare the class label and confidence text
|
||||
std::string classString = classes_[class_id] + std::to_string(confidence).substr(0, 4);
|
||||
|
||||
// Get the size of the text box
|
||||
cv::Size textSize = cv::getTextSize(classString, cv::FONT_HERSHEY_DUPLEX, 0.75, 2, 0);
|
||||
cv::Rect textBox(box.x, box.y - 40, textSize.width + 10, textSize.height + 20);
|
||||
|
||||
// Draw the text box
|
||||
cv::rectangle(frame, textBox, color, cv::FILLED);
|
||||
|
||||
// Put the class label and confidence text above the bounding box
|
||||
cv::putText(frame, classString, cv::Point(box.x + 5, box.y - 10), cv::FONT_HERSHEY_DUPLEX, 0.75, cv::Scalar(0, 0, 0), 2, 0);
|
||||
}
|
||||
} // namespace yolo
|
||||
@@ -0,0 +1,61 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#ifndef YOLO_INFERENCE_H_
|
||||
#define YOLO_INFERENCE_H_
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <openvino/openvino.hpp>
|
||||
|
||||
namespace yolo {
|
||||
|
||||
struct Detection {
|
||||
short class_id;
|
||||
float confidence;
|
||||
cv::Rect box;
|
||||
};
|
||||
|
||||
class Inference {
|
||||
public:
|
||||
Inference() {}
|
||||
// Constructor to initialize the model with default input shape
|
||||
Inference(const std::string &model_path, const float &model_confidence_threshold, const float &model_NMS_threshold);
|
||||
// Constructor to initialize the model with specified input shape
|
||||
Inference(const std::string &model_path, const cv::Size model_input_shape, const float &model_confidence_threshold, const float &model_NMS_threshold);
|
||||
|
||||
void RunInference(cv::Mat &frame);
|
||||
|
||||
private:
|
||||
void InitializeModel(const std::string &model_path);
|
||||
void Preprocessing(const cv::Mat &frame);
|
||||
void PostProcessing(cv::Mat &frame);
|
||||
cv::Rect GetBoundingBox(const cv::Rect &src) const;
|
||||
void DrawDetectedObject(cv::Mat &frame, const Detection &detections) const;
|
||||
|
||||
cv::Point2f scale_factor_; // Scaling factor for the input frame
|
||||
cv::Size2f model_input_shape_; // Input shape of the model
|
||||
cv::Size model_output_shape_; // Output shape of the model
|
||||
|
||||
ov::InferRequest inference_request_; // OpenVINO inference request
|
||||
ov::CompiledModel compiled_model_; // OpenVINO compiled model
|
||||
|
||||
float model_confidence_threshold_; // Confidence threshold for detections
|
||||
float model_NMS_threshold_; // Non-Maximum Suppression threshold
|
||||
|
||||
std::vector<std::string> classes_ {
|
||||
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
|
||||
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
|
||||
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
|
||||
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
|
||||
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
|
||||
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
|
||||
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard",
|
||||
"cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
|
||||
"scissors", "teddy bear", "hair drier", "toothbrush"
|
||||
};
|
||||
};
|
||||
|
||||
} // namespace yolo
|
||||
|
||||
#endif // YOLO_INFERENCE_H_
|
||||
@@ -0,0 +1,42 @@
|
||||
#include "inference.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <opencv2/highgui.hpp>
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
// Check if the correct number of arguments is provided
|
||||
if (argc != 3) {
|
||||
std::cerr << "usage: " << argv[0] << " <model_path> <image_path>" << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Get the model and image paths from the command-line arguments
|
||||
const std::string model_path = argv[1];
|
||||
const std::string image_path = argv[2];
|
||||
|
||||
// Read the input image
|
||||
cv::Mat image = cv::imread(image_path);
|
||||
|
||||
// Check if the image was successfully loaded
|
||||
if (image.empty()) {
|
||||
std::cerr << "ERROR: image is empty" << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Define the confidence and NMS thresholds
|
||||
const float confidence_threshold = 0.5;
|
||||
const float NMS_threshold = 0.5;
|
||||
|
||||
// Initialize the YOLO inference with the specified model and parameters
|
||||
yolo::Inference inference(model_path, cv::Size(640, 640), confidence_threshold, NMS_threshold);
|
||||
|
||||
// Run inference on the input image
|
||||
inference.RunInference(image);
|
||||
|
||||
// Display the image with the detections
|
||||
cv::imshow("image", image);
|
||||
cv::waitKey(0);
|
||||
cv::destroyAllWindows();
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
# Regions Counting Using Ultralytics YOLO11 (Inference on Video)
|
||||
|
||||
> **Note:** Region Counter is now part of **[Ultralytics Solutions](https://docs.ultralytics.com/solutions/)**, offering enhanced features and ongoing updates.
|
||||
>
|
||||
> 🔗 **Explore the official [Region Counting Guide](https://docs.ultralytics.com/guides/region-counting/) for the latest implementation.**
|
||||
|
||||
> 🔔 **Notice:**
|
||||
>
|
||||
> This GitHub example (`ultralytics/examples/YOLOv8-Region-Counter/`) will remain available but **is no longer actively maintained**. For the most current features, updates, and support, please refer to the official [Region Counting guide](https://docs.ultralytics.com/guides/region-counting/) within the Ultralytics documentation. Thank you!
|
||||
|
||||
Region counting is a technique used to count objects within predefined areas or zones in a video feed. This allows for more detailed analysis, especially when monitoring multiple distinct areas simultaneously. Users can interactively adjust these regions by clicking and dragging with the left mouse button, enabling real-time counting tailored to specific needs and layouts. This method is valuable in various [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) applications, from traffic analysis to retail analytics.
|
||||
|
||||
<div>
|
||||
<p align="center">
|
||||
<img src="https://github.com/RizwanMunawar/ultralytics/assets/62513924/5ab3bbd7-fd12-4849-928e-5f294d6c3fcf" width="45%" alt="YOLOv8 region counting visual 1">
|
||||
<img src="https://github.com/RizwanMunawar/ultralytics/assets/62513924/e7c1aea7-474d-4d78-8d48-b50854ffe1ca" width="45%" alt="YOLOv8 region counting visual 2">
|
||||
</p>
|
||||
</div>
|
||||
|
||||
## 📚 Table of Contents
|
||||
|
||||
- [Step 1: Install the Required Libraries](#-step-1-install-the-required-libraries)
|
||||
- [Step 2: Run Region Counting with Ultralytics YOLO11](#-step-2-run-region-counting-with-ultralytics-yolo11)
|
||||
- [Usage Options](#-usage-options)
|
||||
- [Frequently Asked Questions (FAQ)](#-frequently-asked-questions-faq)
|
||||
- [Contributing](#-contributing)
|
||||
|
||||
## ⚙️ Step 1: Install the Required Libraries
|
||||
|
||||
First, clone the Ultralytics repository and navigate to the example directory. Ensure you have Python installed along with necessary dependencies like [PyTorch](https://pytorch.org/) and [OpenCV](https://opencv.org/).
|
||||
|
||||
```bash
|
||||
# Clone the ultralytics repository
|
||||
git clone https://github.com/ultralytics/ultralytics
|
||||
|
||||
# Navigate to the example directory
|
||||
cd ultralytics/examples/YOLOv8-Region-Counter
|
||||
|
||||
# Install required packages (if not already installed)
|
||||
pip install ultralytics shapely
|
||||
```
|
||||
|
||||
## ▶️ Step 2: Run Region Counting with Ultralytics YOLO11
|
||||
|
||||
Execute the script using the following commands. You can customize the source, model, device, and other parameters.
|
||||
|
||||
### Note
|
||||
|
||||
Once the video starts playing, you can dynamically reposition the counting regions within the video frame by clicking and dragging them with your left mouse button.
|
||||
|
||||
```bash
|
||||
# Run inference on a video source, saving results and viewing output
|
||||
python yolov8_region_counter.py --source "path/to/video.mp4" --save-img --view-img
|
||||
|
||||
# Run inference using the CPU
|
||||
python yolov8_region_counter.py --source "path/to/video.mp4" --save-img --view-img --device cpu
|
||||
|
||||
# Use a specific Ultralytics YOLO11 model file
|
||||
python yolov8_region_counter.py --source "path/to/video.mp4" --save-img --weights "path/to/yolo11n.pt"
|
||||
|
||||
# Detect only specific classes (e.g., class 0 and class 2)
|
||||
python yolov8_region_counter.py --source "path/to/video.mp4" --classes 0 2 --weights "path/to/yolo11m.pt"
|
||||
|
||||
# Run inference without saving the output video/images
|
||||
python yolov8_region_counter.py --source "path/to/video.mp4" --view-img
|
||||
```
|
||||
|
||||
Learn more about inference arguments in the Ultralytics [Predict Mode documentation](https://docs.ultralytics.com/modes/predict/).
|
||||
|
||||
## 🛠️ Usage Options
|
||||
|
||||
The script accepts several command-line arguments for customization:
|
||||
|
||||
- `--source`: Path to the input video file.
|
||||
- `--device`: Computation device (`cpu` or GPU ID like `0`).
|
||||
- `--save-img`: Boolean flag to save output frames with detections.
|
||||
- `--view-img`: Boolean flag to display the output video stream in real-time.
|
||||
- `--weights`: Path to the [Ultralytics YOLO11](https://docs.ultralytics.com/models/yolo11/) model file (`.pt`). Defaults typically use a standard model like `yolo11n.pt`.
|
||||
- `--classes`: Filter detections by specific class IDs (e.g., `--classes 0 2 3` to detect classes 0, 2, and 3).
|
||||
- `--line-thickness`: Thickness of the [bounding box](https://www.ultralytics.com/glossary/bounding-box) lines.
|
||||
- `--region-thickness`: Thickness of the lines defining the counting regions.
|
||||
- `--track-thickness`: Thickness of the object tracking lines.
|
||||
|
||||
Explore different models and training options in the [Ultralytics documentation](https://docs.ultralytics.com/).
|
||||
|
||||
## ❓ Frequently Asked Questions (FAQ)
|
||||
|
||||
### What is Region Counting?
|
||||
|
||||
Region counting is a process in [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) used to determine the number of objects within predefined areas of an image or video frame. It's commonly applied in fields like [image processing](https://en.wikipedia.org/wiki/Image_processing) and [pattern recognition](https://en.wikipedia.org/wiki/Pattern_recognition) for analyzing spatial distributions and segmenting objects based on location.
|
||||
|
||||
### Does the Region Counter Support Custom Region Shapes?
|
||||
|
||||
Yes, the Region Counter allows defining regions using polygons (including rectangles). You can customize region properties like coordinates, names, and colors directly in the script. The `shapely` library is used for polygon definitions. See the [Shapely User Manual](https://shapely.readthedocs.io/en/stable/manual.html#polygons) for more details on polygon creation.
|
||||
|
||||
```python
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
# Example definition of counting regions
|
||||
counting_regions = [
|
||||
{
|
||||
"name": "Region 1 (Pentagon)",
|
||||
"polygon": Polygon([(50, 80), (250, 20), (450, 80), (400, 350), (100, 350)]), # 5-point polygon
|
||||
"counts": 0,
|
||||
"dragging": False,
|
||||
"region_color": (255, 42, 4), # BGR color for region
|
||||
"text_color": (255, 255, 255), # BGR color for text
|
||||
},
|
||||
{
|
||||
"name": "Region 2 (Rectangle)",
|
||||
"polygon": Polygon([(200, 250), (440, 250), (440, 550), (200, 550)]), # 4-point polygon (rectangle)
|
||||
"counts": 0,
|
||||
"dragging": False,
|
||||
"region_color": (37, 255, 225), # BGR color for region
|
||||
"text_color": (0, 0, 0), # BGR color for text
|
||||
},
|
||||
]
|
||||
```
|
||||
|
||||
### Why Combine Region Counting with YOLO11?
|
||||
|
||||
[Ultralytics YOLO11](https://docs.ultralytics.com/models/yolo11/) excels at [object detection](https://www.ultralytics.com/glossary/object-detection) and [tracking](https://www.ultralytics.com/glossary/object-tracking) in video streams. Integrating region counting enhances its capabilities by enabling object quantification within specific zones, making it useful for applications like crowd monitoring, traffic flow analysis, and retail footfall counting. Check out our blog post on [Object Detection and Tracking with Ultralytics YOLOv8](https://www.ultralytics.com/blog/object-detection-and-tracking-with-ultralytics-yolov8).
|
||||
|
||||
### How Can I Troubleshoot Issues?
|
||||
|
||||
For debugging, you can enable more verbose output. While this specific script doesn't have a dedicated `--debug` flag, you can add print statements within the code to inspect variables or use standard Python debugging tools. Ensure your video path and model weights path are correct. For common issues, refer to the [Ultralytics FAQ](https://docs.ultralytics.com/help/FAQ/).
|
||||
|
||||
### Can I Use Other YOLO Versions or Custom Models?
|
||||
|
||||
Yes, you can use different Ultralytics YOLO model versions (like YOLOv5, YOLOv9, YOLOv10, YOLO11) or your own custom-trained models by specifying the path to the `.pt` file using the `--weights` argument. Ensure the model is compatible with the Ultralytics framework. Find more about training custom models in the [Model Training guide](https://docs.ultralytics.com/modes/train/).
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions to improve this example or add new features are welcome! Please feel free to submit Pull Requests or open Issues on the main [Ultralytics repository](https://github.com/ultralytics/ultralytics). Remember to check the official [Region Counting guide](https://docs.ultralytics.com/guides/region-counting/) for the latest maintained version.
|
||||
@@ -0,0 +1,253 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from shapely.geometry import Polygon
|
||||
from shapely.geometry.point import Point
|
||||
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.utils.files import increment_path
|
||||
from ultralytics.utils.plotting import Annotator, colors
|
||||
|
||||
track_history = defaultdict(list)
|
||||
|
||||
current_region = None
|
||||
counting_regions = [
|
||||
{
|
||||
"name": "Ultralytics YOLO Polygon Region",
|
||||
"polygon": Polygon([(50, 80), (250, 20), (450, 80), (400, 350), (100, 350)]), # Polygon points
|
||||
"counts": 0,
|
||||
"dragging": False,
|
||||
"region_color": (255, 42, 4), # BGR Value
|
||||
"text_color": (255, 255, 255), # Region Text Color
|
||||
},
|
||||
{
|
||||
"name": "Ultralytics YOLO Rectangle Region",
|
||||
"polygon": Polygon([(200, 250), (440, 250), (440, 550), (200, 550)]), # Polygon points
|
||||
"counts": 0,
|
||||
"dragging": False,
|
||||
"region_color": (37, 255, 225), # BGR Value
|
||||
"text_color": (0, 0, 0), # Region Text Color
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def mouse_callback(event: int, x: int, y: int, flags: int, param: Any) -> None:
|
||||
"""Handle mouse events for region manipulation in the video frame.
|
||||
|
||||
This function enables interactive region selection and dragging functionality for counting regions. It responds to
|
||||
mouse button down, move, and up events to allow users to select and reposition counting regions in real-time.
|
||||
|
||||
Args:
|
||||
event (int): The mouse event type (e.g., cv2.EVENT_LBUTTONDOWN, cv2.EVENT_MOUSEMOVE).
|
||||
x (int): The x-coordinate of the mouse pointer.
|
||||
y (int): The y-coordinate of the mouse pointer.
|
||||
flags (int): Additional flags passed by OpenCV.
|
||||
param (Any): Additional parameters passed to the callback.
|
||||
|
||||
Examples:
|
||||
Set up mouse callback for interactive region manipulation
|
||||
>>> cv2.setMouseCallback("window_name", mouse_callback)
|
||||
"""
|
||||
global current_region
|
||||
|
||||
# Mouse left button down event
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
for region in counting_regions:
|
||||
if region["polygon"].contains(Point((x, y))):
|
||||
current_region = region
|
||||
current_region["dragging"] = True
|
||||
current_region["offset_x"] = x
|
||||
current_region["offset_y"] = y
|
||||
|
||||
# Mouse move event
|
||||
elif event == cv2.EVENT_MOUSEMOVE:
|
||||
if current_region is not None and current_region["dragging"]:
|
||||
dx = x - current_region["offset_x"]
|
||||
dy = y - current_region["offset_y"]
|
||||
current_region["polygon"] = Polygon(
|
||||
[(p[0] + dx, p[1] + dy) for p in current_region["polygon"].exterior.coords]
|
||||
)
|
||||
current_region["offset_x"] = x
|
||||
current_region["offset_y"] = y
|
||||
|
||||
# Mouse left button up event
|
||||
elif event == cv2.EVENT_LBUTTONUP:
|
||||
if current_region is not None and current_region["dragging"]:
|
||||
current_region["dragging"] = False
|
||||
|
||||
|
||||
def run(
|
||||
weights: str = "yolo11n.pt",
|
||||
source: str | None = None,
|
||||
device: str = "cpu",
|
||||
view_img: bool = False,
|
||||
save_img: bool = False,
|
||||
exist_ok: bool = False,
|
||||
classes: list[int] | None = None,
|
||||
line_thickness: int = 2,
|
||||
track_thickness: int = 2,
|
||||
region_thickness: int = 2,
|
||||
) -> None:
|
||||
"""Run object detection and counting within specified regions using YOLO and ByteTrack.
|
||||
|
||||
This function performs real-time object detection, tracking, and counting within user-defined polygonal or
|
||||
rectangular regions. It supports interactive region manipulation, multiple counting areas, and both live viewing and
|
||||
video saving capabilities.
|
||||
|
||||
Args:
|
||||
weights (str): Path to the YOLO model weights file.
|
||||
source (str): Path to the input video file.
|
||||
device (str): Processing device specification ('cpu', '0', '1', etc.).
|
||||
view_img (bool): Display results in a live window.
|
||||
save_img (bool): Save processed video to file.
|
||||
exist_ok (bool): Overwrite existing output files without incrementing.
|
||||
classes (list[int], optional): Specific class IDs to detect and track.
|
||||
line_thickness (int): Thickness of bounding box lines.
|
||||
track_thickness (int): Thickness of object tracking lines.
|
||||
region_thickness (int): Thickness of counting region boundaries.
|
||||
|
||||
Examples:
|
||||
Run region counting with default settings
|
||||
>>> run(source="video.mp4", view_img=True)
|
||||
|
||||
Run with custom model and specific classes
|
||||
>>> run(weights="yolo11s.pt", source="traffic.mp4", classes=[0, 2, 3], device="0")
|
||||
"""
|
||||
vid_frame_count = 0
|
||||
|
||||
# Check source path
|
||||
if not Path(source).exists():
|
||||
raise FileNotFoundError(f"Source path '{source}' does not exist.")
|
||||
|
||||
# Setup Model
|
||||
model = YOLO(f"{weights}")
|
||||
model.to("cuda") if device == "0" else model.to("cpu")
|
||||
|
||||
# Extract classes names
|
||||
names = model.names
|
||||
|
||||
# Video setup
|
||||
videocapture = cv2.VideoCapture(source)
|
||||
frame_width = int(videocapture.get(3))
|
||||
frame_height = int(videocapture.get(4))
|
||||
fps = int(videocapture.get(5))
|
||||
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
|
||||
# Output setup
|
||||
save_dir = increment_path(Path("ultralytics_rc_output") / "exp", exist_ok)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
video_writer = cv2.VideoWriter(str(save_dir / f"{Path(source).stem}.avi"), fourcc, fps, (frame_width, frame_height))
|
||||
|
||||
# Iterate over video frames
|
||||
while videocapture.isOpened():
|
||||
success, frame = videocapture.read()
|
||||
if not success:
|
||||
break
|
||||
vid_frame_count += 1
|
||||
|
||||
# Extract the results
|
||||
results = model.track(frame, persist=True, classes=classes)
|
||||
|
||||
if results[0].boxes.is_track:
|
||||
boxes = results[0].boxes.xyxy.cpu()
|
||||
track_ids = results[0].boxes.id.int().cpu().tolist()
|
||||
clss = results[0].boxes.cls.cpu().tolist()
|
||||
|
||||
annotator = Annotator(frame, line_width=line_thickness, example=str(names))
|
||||
|
||||
for box, track_id, cls in zip(boxes, track_ids, clss):
|
||||
annotator.box_label(box, str(names[cls]), color=colors(cls, True))
|
||||
bbox_center = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2 # Bbox center
|
||||
|
||||
track = track_history[track_id] # Tracking Lines plot
|
||||
track.append((float(bbox_center[0]), float(bbox_center[1])))
|
||||
if len(track) > 30:
|
||||
track.pop(0)
|
||||
points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
|
||||
cv2.polylines(frame, [points], isClosed=False, color=colors(cls, True), thickness=track_thickness)
|
||||
|
||||
# Check if detection inside region
|
||||
for region in counting_regions:
|
||||
if region["polygon"].contains(Point((bbox_center[0], bbox_center[1]))):
|
||||
region["counts"] += 1
|
||||
|
||||
# Draw regions (Polygons/Rectangles)
|
||||
for region in counting_regions:
|
||||
region_label = str(region["counts"])
|
||||
region_color = region["region_color"]
|
||||
region_text_color = region["text_color"]
|
||||
|
||||
polygon_coordinates = np.array(region["polygon"].exterior.coords, dtype=np.int32)
|
||||
centroid_x, centroid_y = int(region["polygon"].centroid.x), int(region["polygon"].centroid.y)
|
||||
|
||||
text_size, _ = cv2.getTextSize(
|
||||
region_label, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.7, thickness=line_thickness
|
||||
)
|
||||
text_x = centroid_x - text_size[0] // 2
|
||||
text_y = centroid_y + text_size[1] // 2
|
||||
cv2.rectangle(
|
||||
frame,
|
||||
(text_x - 5, text_y - text_size[1] - 5),
|
||||
(text_x + text_size[0] + 5, text_y + 5),
|
||||
region_color,
|
||||
-1,
|
||||
)
|
||||
cv2.putText(
|
||||
frame, region_label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, region_text_color, line_thickness
|
||||
)
|
||||
cv2.polylines(frame, [polygon_coordinates], isClosed=True, color=region_color, thickness=region_thickness)
|
||||
|
||||
if view_img:
|
||||
if vid_frame_count == 1:
|
||||
cv2.namedWindow("Ultralytics YOLO Region Counter Movable")
|
||||
cv2.setMouseCallback("Ultralytics YOLO Region Counter Movable", mouse_callback)
|
||||
cv2.imshow("Ultralytics YOLO Region Counter Movable", frame)
|
||||
|
||||
if save_img:
|
||||
video_writer.write(frame)
|
||||
|
||||
for region in counting_regions: # Reinitialize count for each region
|
||||
region["counts"] = 0
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
del vid_frame_count
|
||||
video_writer.release()
|
||||
videocapture.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def parse_opt() -> argparse.Namespace:
|
||||
"""Parse command line arguments for the region counting application."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--weights", type=str, default="yolo11n.pt", help="initial weights path")
|
||||
parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
|
||||
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("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
|
||||
parser.add_argument("--line-thickness", type=int, default=2, help="bounding box thickness")
|
||||
parser.add_argument("--track-thickness", type=int, default=2, help="Tracking line thickness")
|
||||
parser.add_argument("--region-thickness", type=int, default=4, help="Region thickness")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(options: argparse.Namespace) -> None:
|
||||
"""Execute the main region counting functionality with the provided options."""
|
||||
run(**vars(options))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
opt = parse_opt()
|
||||
main(opt)
|
||||
@@ -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.
|
||||
@@ -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): Flag to show or hide confidences 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()))
|
||||
@@ -0,0 +1,64 @@
|
||||
# YOLOv8-Segmentation-ONNXRuntime-Python Demo
|
||||
|
||||
This repository provides a [Python](https://www.python.org/) demo for performing instance segmentation with [Ultralytics YOLOv8](https://docs.ultralytics.com/models/yolov8/) using [ONNX Runtime](https://onnxruntime.ai/). It highlights the interoperability of YOLOv8 models, allowing inference without requiring the full [PyTorch](https://pytorch.org/) stack. This approach is ideal for deployment scenarios where minimal dependencies are preferred. Learn more about the [segmentation task](https://docs.ultralytics.com/tasks/segment/) on our documentation.
|
||||
|
||||
## ✨ Features
|
||||
|
||||
- **Framework Agnostic**: Runs segmentation inference purely on ONNX Runtime without importing PyTorch.
|
||||
- **Efficient Inference**: Supports both FP32 and [half-precision](https://www.ultralytics.com/glossary/half-precision) (FP16) for [ONNX](https://onnx.ai/) models, catering to different computational needs and optimizing [inference latency](https://www.ultralytics.com/glossary/inference-latency).
|
||||
- **Ease of Use**: Utilizes simple command-line arguments for straightforward model execution.
|
||||
- **Broad Compatibility**: Leverages [NumPy](https://numpy.org/) and [OpenCV](https://opencv.org/) for image processing, ensuring wide compatibility across various environments.
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
Install the required packages using pip. You will need [`ultralytics`](https://github.com/ultralytics/ultralytics) for exporting the YOLOv8-seg ONNX model and using some utility functions, [`onnxruntime-gpu`](https://pypi.org/project/onnxruntime-gpu/) for GPU-accelerated inference, and [`opencv-python`](https://pypi.org/project/opencv-python/) for image processing.
|
||||
|
||||
```bash
|
||||
pip install ultralytics
|
||||
pip install onnxruntime-gpu # For GPU support
|
||||
# pip install onnxruntime # For CPU-only support
|
||||
pip install numpy opencv-python
|
||||
```
|
||||
|
||||
## 🚀 Getting Started
|
||||
|
||||
### 1. Export the YOLOv8 ONNX Model
|
||||
|
||||
First, export your Ultralytics YOLOv8 segmentation model to the ONNX format using the `ultralytics` package. This step converts the PyTorch model into a standardized format suitable for ONNX Runtime. Check our [Export documentation](https://docs.ultralytics.com/modes/export/) for more details on export options and our [ONNX integration guide](https://docs.ultralytics.com/integrations/onnx/).
|
||||
|
||||
```bash
|
||||
yolo export model=yolov8s-seg.pt imgsz=640 format=onnx opset=12 simplify
|
||||
```
|
||||
|
||||
### 2. Run Inference
|
||||
|
||||
Perform inference with the exported ONNX model on your images or video sources. Specify the path to your ONNX model and the image source using the command-line arguments.
|
||||
|
||||
```bash
|
||||
python main.py --model yolov8s-seg.onnx --source path/to/image.jpg
|
||||
```
|
||||
|
||||
### Example Output
|
||||
|
||||
After running the command, the script will process the image, perform segmentation, and display the results with bounding boxes and masks overlaid.
|
||||
|
||||
<img src="https://user-images.githubusercontent.com/51357717/279988626-eb74823f-1563-4d58-a8e4-0494025b7c9a.jpg" alt="YOLOv8 Segmentation ONNX Demo Output" width="800">
|
||||
|
||||
## 💡 Advanced Usage
|
||||
|
||||
For more advanced usage scenarios, such as processing video streams or adjusting inference parameters, please refer to the command-line arguments available in the `main.py` script. You can explore options like confidence thresholds and input image sizes.
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
We welcome contributions to improve this demo! If you encounter bugs, have feature requests, or want to submit enhancements (like a new algorithm or improved processing steps), please open an issue or pull request on the main [Ultralytics repository](https://github.com/ultralytics/ultralytics). See our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for more details on how to get involved.
|
||||
|
||||
## 📄 License
|
||||
|
||||
This project is licensed under the AGPL-3.0 License. For detailed information, please see the [LICENSE](https://github.com/ultralytics/ultralytics/blob/main/LICENSE) file or read the full [AGPL-3.0 license text](https://opensource.org/license/agpl-v3).
|
||||
|
||||
## 🙏 Acknowledgments
|
||||
|
||||
- This YOLOv8-Segmentation-ONNXRuntime-Python demo was contributed by GitHub user [jamjamjon](https://github.com/jamjamjon).
|
||||
- Thanks to the [ONNX Runtime community](https://github.com/microsoft/onnxruntime) for providing a robust and efficient inference engine.
|
||||
|
||||
We hope you find this demo useful! Feel free to contribute and help make it even better.
|
||||
@@ -0,0 +1,179 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import torch
|
||||
|
||||
from ultralytics.engine.results import Results
|
||||
from ultralytics.utils import ASSETS, YAML, nms, ops
|
||||
from ultralytics.utils.checks import check_yaml
|
||||
|
||||
|
||||
class YOLOv8Seg:
|
||||
"""YOLOv8 segmentation model for performing instance segmentation using ONNX Runtime.
|
||||
|
||||
This class implements a YOLOv8 instance segmentation model using ONNX Runtime for inference. It handles
|
||||
preprocessing of input images, running inference with the ONNX model, and postprocessing the results to generate
|
||||
bounding boxes and segmentation masks.
|
||||
|
||||
Attributes:
|
||||
session (ort.InferenceSession): ONNX Runtime inference session for model execution.
|
||||
imgsz (tuple[int, int]): Input image size as (height, width) for the model.
|
||||
classes (dict): Dictionary mapping class indices to class names from the dataset.
|
||||
conf (float): Confidence threshold for filtering detections.
|
||||
iou (float): IoU threshold used by non-maximum suppression.
|
||||
|
||||
Methods:
|
||||
letterbox: Resize and pad image while maintaining aspect ratio.
|
||||
preprocess: Preprocess the input image before feeding it into the model.
|
||||
postprocess: Post-process model predictions to extract meaningful results.
|
||||
process_mask: Process prototype masks with predicted mask coefficients to generate instance segmentation masks.
|
||||
|
||||
Examples:
|
||||
>>> model = YOLOv8Seg("yolov8n-seg.onnx", conf=0.25, iou=0.7)
|
||||
>>> img = cv2.imread("image.jpg")
|
||||
>>> results = model(img)
|
||||
>>> cv2.imshow("Segmentation", results[0].plot())
|
||||
"""
|
||||
|
||||
def __init__(self, onnx_model: str, conf: float = 0.25, iou: float = 0.7, imgsz: int | tuple[int, int] = 640):
|
||||
"""Initialize the instance segmentation model using an ONNX model.
|
||||
|
||||
Args:
|
||||
onnx_model (str): Path to the ONNX model file.
|
||||
conf (float, optional): Confidence threshold for filtering detections.
|
||||
iou (float, optional): IoU threshold for non-maximum suppression.
|
||||
imgsz (int | tuple[int, int], optional): Input image size of the model. Can be an integer for square input
|
||||
or a tuple for rectangular input.
|
||||
"""
|
||||
available = ort.get_available_providers()
|
||||
providers = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in available]
|
||||
self.session = ort.InferenceSession(onnx_model, providers=providers or available)
|
||||
|
||||
self.imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz
|
||||
self.classes = YAML.load(check_yaml("coco8.yaml"))["names"]
|
||||
self.conf = conf
|
||||
self.iou = iou
|
||||
|
||||
def __call__(self, img: np.ndarray) -> list[Results]:
|
||||
"""Run inference on the input image using the ONNX model.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): The original input image in BGR format.
|
||||
|
||||
Returns:
|
||||
(list[Results]): Processed detection results after post-processing, containing bounding boxes and
|
||||
segmentation masks.
|
||||
"""
|
||||
prep_img = self.preprocess(img, self.imgsz)
|
||||
outs = self.session.run(None, {self.session.get_inputs()[0].name: prep_img})
|
||||
return self.postprocess(img, prep_img, outs)
|
||||
|
||||
def letterbox(self, img: np.ndarray, new_shape: tuple[int, int] = (640, 640)) -> np.ndarray:
|
||||
"""Resize and pad image while maintaining aspect ratio.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): Input image in BGR format.
|
||||
new_shape (tuple[int, int], optional): Target shape as (height, width).
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Resized and padded image.
|
||||
"""
|
||||
shape = img.shape[:2] # current shape [height, width]
|
||||
|
||||
# Scale ratio (new / old)
|
||||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||
|
||||
# Compute padding
|
||||
new_unpad = round(shape[1] * r), round(shape[0] * r)
|
||||
dw, dh = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2 # wh padding
|
||||
|
||||
if shape[::-1] != new_unpad: # resize
|
||||
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||
top, bottom = round(dh - 0.1), round(dh + 0.1)
|
||||
left, right = round(dw - 0.1), round(dw + 0.1)
|
||||
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
|
||||
|
||||
return img
|
||||
|
||||
def preprocess(self, img: np.ndarray, new_shape: tuple[int, int]) -> np.ndarray:
|
||||
"""Preprocess the input image before feeding it into the model.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): The input image in BGR format.
|
||||
new_shape (tuple[int, int]): The target shape for resizing as (height, width).
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Preprocessed image ready for model inference, with shape (1, 3, height, width) and normalized
|
||||
to [0, 1].
|
||||
"""
|
||||
img = self.letterbox(img, new_shape)
|
||||
img = img[..., ::-1].transpose([2, 0, 1])[None] # BGR to RGB, BHWC to BCHW
|
||||
img = np.ascontiguousarray(img)
|
||||
img = img.astype(np.float32) / 255 # Normalize to [0, 1]
|
||||
return img
|
||||
|
||||
def postprocess(self, img: np.ndarray, prep_img: np.ndarray, outs: list) -> list[Results]:
|
||||
"""Post-process model predictions to extract meaningful results.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): The original input image.
|
||||
prep_img (np.ndarray): The preprocessed image used for inference.
|
||||
outs (list): Model outputs containing predictions and prototype masks.
|
||||
|
||||
Returns:
|
||||
(list[Results]): Processed detection results containing bounding boxes and segmentation masks.
|
||||
"""
|
||||
preds, protos = (torch.from_numpy(p) for p in outs)
|
||||
preds = nms.non_max_suppression(preds, self.conf, self.iou, nc=len(self.classes))
|
||||
|
||||
results = []
|
||||
for i, pred in enumerate(preds):
|
||||
pred[:, :4] = ops.scale_boxes(prep_img.shape[2:], pred[:, :4], img.shape)
|
||||
masks = self.process_mask(protos[i], pred[:, 6:], pred[:, :4], img.shape[:2])
|
||||
results.append(Results(img, path="", names=self.classes, boxes=pred[:, :6], masks=masks))
|
||||
|
||||
return results
|
||||
|
||||
def process_mask(
|
||||
self, protos: torch.Tensor, masks_in: torch.Tensor, bboxes: torch.Tensor, shape: tuple[int, int]
|
||||
) -> torch.Tensor:
|
||||
"""Process prototype masks with predicted mask coefficients to generate instance segmentation masks.
|
||||
|
||||
Args:
|
||||
protos (torch.Tensor): Prototype masks with shape (mask_dim, mask_h, mask_w).
|
||||
masks_in (torch.Tensor): Predicted mask coefficients with shape (N, mask_dim), where N is number of
|
||||
detections.
|
||||
bboxes (torch.Tensor): Bounding boxes with shape (N, 4), where N is number of detections.
|
||||
shape (tuple[int, int]): The size of the input image as (height, width).
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): Binary segmentation masks with shape (N, height, width).
|
||||
"""
|
||||
c, mh, mw = protos.shape # CHW
|
||||
masks = (masks_in @ protos.float().view(c, -1)).view(-1, mh, mw) # Matrix multiplication
|
||||
masks = ops.scale_masks(masks[None], shape)[0] # Scale masks to original image size
|
||||
masks = ops.crop_mask(masks, bboxes) # Crop masks to bounding boxes
|
||||
return masks.gt_(0.0) # Convert to binary masks
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", type=str, required=True, help="Path to ONNX model")
|
||||
parser.add_argument("--source", type=str, default=str(ASSETS / "bus.jpg"), help="Path to input image")
|
||||
parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
|
||||
parser.add_argument("--iou", type=float, default=0.7, help="NMS IoU threshold")
|
||||
args = parser.parse_args()
|
||||
|
||||
model = YOLOv8Seg(args.model, args.conf, args.iou)
|
||||
img = cv2.imread(args.source)
|
||||
results = model(img)
|
||||
|
||||
cv2.imshow("Segmented Image", results[0].plot())
|
||||
cv2.waitKey(0)
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,79 @@
|
||||
# Running Ultralytics YOLOv8 with TFLite Runtime
|
||||
|
||||
This guide demonstrates how to perform inference using an Ultralytics [YOLOv8](https://docs.ultralytics.com/models/yolov8/) model exported to the [TensorFlow Lite (TFLite)](https://ai.google.dev/edge/litert) format. TFLite is a popular choice for deploying machine learning models on mobile, embedded, and IoT devices due to its optimization for on-device inference with low latency and a small binary size. This example supports FP32, FP16, and INT8 quantized TFLite models.
|
||||
|
||||
## ⚙️ Installation
|
||||
|
||||
Before running inference, you need to install the necessary TFLite interpreter package. Choose the appropriate package based on your hardware (CPU or GPU).
|
||||
|
||||
### Installing `tflite-runtime` (Recommended for Edge Devices)
|
||||
|
||||
The `tflite-runtime` package is a smaller package that includes the bare minimum required to run inferences with TensorFlow Lite, primarily the `Interpreter` Python class. It's ideal for resource-constrained environments like Raspberry Pi or Coral Edge TPU.
|
||||
|
||||
```bash
|
||||
pip install tflite-runtime
|
||||
```
|
||||
|
||||
For more details, refer to the official [TFLite Python quickstart guide](https://ai.google.dev/edge/litert/microcontrollers/python).
|
||||
|
||||
### Installing Full `tensorflow` Package (CPU or GPU)
|
||||
|
||||
Alternatively, you can install the full TensorFlow package. This includes the TFLite interpreter along with the complete TensorFlow library.
|
||||
|
||||
- **CPU-Only:** Suitable if you don't have an NVIDIA GPU or don't need GPU acceleration.
|
||||
|
||||
```bash
|
||||
pip install tensorflow
|
||||
```
|
||||
|
||||
- **GPU Support:** To leverage NVIDIA GPU acceleration for potentially faster inference, install `tensorflow` with GPU support. Ensure you have the necessary [NVIDIA drivers](https://www.nvidia.com/Download/index.aspx) and CUDA toolkit installed.
|
||||
|
||||
```bash
|
||||
# Check TensorFlow documentation for specific CUDA/cuDNN version requirements
|
||||
pip install tensorflow[and-cuda] # Or follow specific instructions on TF website
|
||||
```
|
||||
|
||||
Visit the official [TensorFlow installation guide](https://www.tensorflow.org/install) for detailed instructions, including GPU setup.
|
||||
|
||||
## 🚀 Usage
|
||||
|
||||
Follow these steps to run inference with your exported YOLOv8 TFLite model.
|
||||
|
||||
1. **Export YOLOv8 Model to TFLite:**
|
||||
First, export your trained Ultralytics YOLOv8 model (e.g., `yolov8n.pt`) to the TFLite format using the `yolo export` command. This example exports an INT8 quantized model for optimal performance on edge devices. You can also export FP32 or FP16 models by adjusting the `format` and quantization arguments. Refer to the Ultralytics [Export mode documentation](https://docs.ultralytics.com/modes/export/) for more options.
|
||||
|
||||
```bash
|
||||
yolo export model=yolov8n.pt imgsz=640 format=tflite int8=True # Exports yolov8n_saved_model/yolov8n_full_integer_quant.tflite
|
||||
```
|
||||
|
||||
The export process will create a directory (e.g., `yolov8n_saved_model`) containing the `.tflite` model file and potentially a `metadata.yaml` file with class names and other model details.
|
||||
|
||||
2. **Run Inference Script:**
|
||||
Execute the provided Python script (`main.py`) to perform inference on an image. Adjust the arguments as needed for your specific model path, image source, confidence threshold, and IoU threshold.
|
||||
|
||||
```bash
|
||||
python main.py \
|
||||
--model yolov8n_saved_model/yolov8n_full_integer_quant.tflite \
|
||||
--img image.jpg \
|
||||
--conf 0.25 \
|
||||
--iou 0.45 \
|
||||
--metadata yolov8n_saved_model/metadata.yaml
|
||||
```
|
||||
|
||||
- `--model`: Path to the exported `.tflite` model file.
|
||||
- `--img`: Path to the input image for detection.
|
||||
- `--conf`: Minimum [confidence threshold](https://www.ultralytics.com/glossary/confidence) for detections (e.g., 0.25).
|
||||
- `--iou`: [Intersection over Union (IoU)](https://www.ultralytics.com/glossary/intersection-over-union-iou) threshold for Non-Maximum Suppression (NMS).
|
||||
- `--metadata`: Path to the `metadata.yaml` file generated during export (contains class names).
|
||||
|
||||
## ✅ Output
|
||||
|
||||
The script will process the input image using the specified TFLite model and display the image with bounding boxes drawn around detected objects. Each box will be labeled with the predicted class name and confidence score.
|
||||
|
||||
|
||||
|
||||
This example provides a straightforward way to deploy Ultralytics YOLOv8 models on devices supporting TFLite, enabling efficient **object detection** in various applications. Explore different [quantization](https://www.ultralytics.com/glossary/model-quantization) options and model sizes to find the best balance between performance and accuracy for your specific use case.
|
||||
|
||||
## 🤝 Contribute
|
||||
|
||||
Contributions to enhance this example or add new functionalities are welcome! Feel free to fork the [Ultralytics repository](https://github.com/ultralytics/ultralytics), make your changes, and submit a pull request.
|
||||
260
algorithms/dms_yolo/code/examples/YOLOv8-TFLite-Python/main.py
Normal file
260
algorithms/dms_yolo/code/examples/YOLOv8-TFLite-Python/main.py
Normal file
@@ -0,0 +1,260 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
from ultralytics.utils import ASSETS
|
||||
|
||||
try:
|
||||
from tflite_runtime.interpreter import Interpreter
|
||||
except ImportError:
|
||||
import tensorflow as tf
|
||||
|
||||
Interpreter = tf.lite.Interpreter
|
||||
|
||||
|
||||
class YOLOv8TFLite:
|
||||
"""A YOLOv8 object detection class using TensorFlow Lite for efficient inference.
|
||||
|
||||
This class handles model loading, preprocessing, inference, and visualization of detection results for YOLOv8 models
|
||||
converted to TensorFlow Lite format.
|
||||
|
||||
Attributes:
|
||||
model (Interpreter): TensorFlow Lite interpreter for the YOLOv8 model.
|
||||
conf (float): Confidence threshold for filtering detections.
|
||||
iou (float): Intersection over Union threshold for non-maximum suppression.
|
||||
classes (dict): Dictionary mapping class IDs to class names.
|
||||
color_palette (np.ndarray): Random color palette for visualization with shape (num_classes, 3).
|
||||
in_width (int): Input width required by the model.
|
||||
in_height (int): Input height required by the model.
|
||||
in_index (int): Input tensor index in the model.
|
||||
in_scale (float): Input quantization scale factor.
|
||||
in_zero_point (int): Input quantization zero point.
|
||||
int8 (bool): Whether the model uses int8 quantization.
|
||||
out_index (int): Output tensor index in the model.
|
||||
out_scale (float): Output quantization scale factor.
|
||||
out_zero_point (int): Output quantization zero point.
|
||||
|
||||
Methods:
|
||||
letterbox: Resize and pad image while maintaining aspect ratio.
|
||||
draw_detections: Draw bounding boxes and labels on the input image.
|
||||
preprocess: Preprocess the input image before inference.
|
||||
postprocess: Process model outputs to extract and visualize detections.
|
||||
detect: Perform object detection on an input image.
|
||||
|
||||
Examples:
|
||||
Initialize detector and run inference
|
||||
>>> detector = YOLOv8TFLite("yolov8n.tflite", conf=0.25, iou=0.45)
|
||||
>>> result = detector.detect("image.jpg")
|
||||
>>> cv2.imshow("Result", result)
|
||||
"""
|
||||
|
||||
def __init__(self, model: str, conf: float = 0.25, iou: float = 0.45, metadata: str | None = None):
|
||||
"""Initialize the YOLOv8TFLite detector.
|
||||
|
||||
Args:
|
||||
model (str): Path to the TFLite model file.
|
||||
conf (float): Confidence threshold for filtering detections.
|
||||
iou (float): IoU threshold for non-maximum suppression.
|
||||
metadata (str | None): Path to the metadata file containing class names.
|
||||
"""
|
||||
self.conf = conf
|
||||
self.iou = iou
|
||||
if metadata is None:
|
||||
self.classes = {i: i for i in range(1000)}
|
||||
else:
|
||||
with open(metadata) as f:
|
||||
self.classes = yaml.safe_load(f)["names"]
|
||||
np.random.seed(42) # Set seed for reproducible colors
|
||||
self.color_palette = np.random.uniform(128, 255, size=(len(self.classes), 3))
|
||||
|
||||
# Initialize the TFLite interpreter
|
||||
self.model = Interpreter(model_path=model)
|
||||
self.model.allocate_tensors()
|
||||
|
||||
# Get input details
|
||||
input_details = self.model.get_input_details()[0]
|
||||
self.in_width, self.in_height = input_details["shape"][1:3]
|
||||
self.in_index = input_details["index"]
|
||||
self.in_scale, self.in_zero_point = input_details["quantization"]
|
||||
self.int8 = input_details["dtype"] == np.int8
|
||||
|
||||
# Get output details
|
||||
output_details = self.model.get_output_details()[0]
|
||||
self.out_index = output_details["index"]
|
||||
self.out_scale, self.out_zero_point = output_details["quantization"]
|
||||
|
||||
def letterbox(
|
||||
self, img: np.ndarray, new_shape: tuple[int, int] = (640, 640)
|
||||
) -> tuple[np.ndarray, tuple[float, float]]:
|
||||
"""Resize and pad image while maintaining aspect ratio.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): Input image with shape (H, W, C).
|
||||
new_shape (tuple[int, int]): Target shape (height, width).
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Resized and padded image.
|
||||
(tuple[float, float]): Padding ratios (top/height, left/width) for coordinate adjustment.
|
||||
"""
|
||||
shape = img.shape[:2] # Current shape [height, width]
|
||||
|
||||
# Scale ratio (new / old)
|
||||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||
|
||||
# Compute padding
|
||||
new_unpad = round(shape[1] * r), round(shape[0] * r)
|
||||
dw, dh = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2 # wh padding
|
||||
|
||||
if shape[::-1] != new_unpad: # Resize if needed
|
||||
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||
top, bottom = round(dh - 0.1), round(dh + 0.1)
|
||||
left, right = round(dw - 0.1), round(dw + 0.1)
|
||||
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
|
||||
|
||||
return img, (top / img.shape[0], left / img.shape[1])
|
||||
|
||||
def draw_detections(self, img: np.ndarray, box: np.ndarray, score: np.float32, class_id: int) -> None:
|
||||
"""Draw bounding boxes and labels on the input image based on detected objects.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): The input image to draw detections on.
|
||||
box (np.ndarray): Detected bounding box in the format [x1, y1, width, height].
|
||||
score (np.float32): Confidence score of the detection.
|
||||
class_id (int): Class ID for the detected object.
|
||||
"""
|
||||
x1, y1, w, h = box
|
||||
color = self.color_palette[class_id]
|
||||
|
||||
# Draw bounding box
|
||||
cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)
|
||||
|
||||
# Create label with class name and score
|
||||
label = f"{self.classes[class_id]}: {score:.2f}"
|
||||
|
||||
# Get text size for background rectangle
|
||||
(label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||||
|
||||
# Position label above or below box depending on space
|
||||
label_x = x1
|
||||
label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10
|
||||
|
||||
# Draw label background
|
||||
cv2.rectangle(
|
||||
img,
|
||||
(int(label_x), int(label_y - label_height)),
|
||||
(int(label_x + label_width), int(label_y + label_height)),
|
||||
color,
|
||||
cv2.FILLED,
|
||||
)
|
||||
|
||||
# Draw text
|
||||
cv2.putText(img, label, (int(label_x), int(label_y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
|
||||
|
||||
def preprocess(self, img: np.ndarray) -> tuple[np.ndarray, tuple[float, float]]:
|
||||
"""Preprocess the input image before performing inference.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): The input image to be preprocessed with shape (H, W, C).
|
||||
|
||||
Returns:
|
||||
(np.ndarray): Preprocessed image ready for model input.
|
||||
(tuple[float, float]): Padding ratios for coordinate adjustment.
|
||||
"""
|
||||
img, pad = self.letterbox(img, (self.in_width, self.in_height))
|
||||
img = img[..., ::-1][None] # BGR to RGB and add batch dimension (N, H, W, C) for TFLite
|
||||
img = np.ascontiguousarray(img)
|
||||
img = img.astype(np.float32)
|
||||
return img / 255, pad # Normalize to [0, 1]
|
||||
|
||||
def postprocess(self, img: np.ndarray, outputs: np.ndarray, pad: tuple[float, float]) -> np.ndarray:
|
||||
"""Process model outputs to extract and visualize detections.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): The original input image.
|
||||
outputs (np.ndarray): Raw model outputs.
|
||||
pad (tuple[float, float]): Padding ratios from preprocessing.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The input image with detections drawn on it.
|
||||
"""
|
||||
# Adjust coordinates based on padding and scale to original image size
|
||||
outputs[:, 0] -= pad[1]
|
||||
outputs[:, 1] -= pad[0]
|
||||
outputs[:, :4] *= max(img.shape)
|
||||
|
||||
# Transform outputs to [x, y, w, h] format
|
||||
outputs = outputs.transpose(0, 2, 1)
|
||||
outputs[..., 0] -= outputs[..., 2] / 2 # x center to top-left x
|
||||
outputs[..., 1] -= outputs[..., 3] / 2 # y center to top-left y
|
||||
|
||||
for out in outputs:
|
||||
# Get scores and apply confidence threshold
|
||||
scores = out[:, 4:].max(-1)
|
||||
keep = scores > self.conf
|
||||
boxes = out[keep, :4]
|
||||
scores = scores[keep]
|
||||
class_ids = out[keep, 4:].argmax(-1)
|
||||
|
||||
# Apply non-maximum suppression
|
||||
indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf, self.iou).flatten()
|
||||
|
||||
# Draw detections that survived NMS
|
||||
[self.draw_detections(img, boxes[i], scores[i], class_ids[i]) for i in indices]
|
||||
|
||||
return img
|
||||
|
||||
def detect(self, img_path: str) -> np.ndarray:
|
||||
"""Perform object detection on an input image.
|
||||
|
||||
Args:
|
||||
img_path (str): Path to the input image file.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The output image with drawn detections.
|
||||
"""
|
||||
# Load and preprocess image
|
||||
img = cv2.imread(img_path)
|
||||
x, pad = self.preprocess(img)
|
||||
|
||||
# Apply quantization if model is int8
|
||||
if self.int8:
|
||||
x = (x / self.in_scale + self.in_zero_point).astype(np.int8)
|
||||
|
||||
# Set input tensor and run inference
|
||||
self.model.set_tensor(self.in_index, x)
|
||||
self.model.invoke()
|
||||
|
||||
# Get output and dequantize if necessary
|
||||
y = self.model.get_tensor(self.out_index)
|
||||
if self.int8:
|
||||
y = (y.astype(np.float32) - self.out_zero_point) * self.out_scale
|
||||
|
||||
# Process detections and return result
|
||||
return self.postprocess(img, y, pad)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="yolov8n_saved_model/yolov8n_full_integer_quant.tflite",
|
||||
help="Path to TFLite model.",
|
||||
)
|
||||
parser.add_argument("--img", type=str, default=str(ASSETS / "bus.jpg"), help="Path to input image")
|
||||
parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
|
||||
parser.add_argument("--iou", type=float, default=0.45, help="NMS IoU threshold")
|
||||
parser.add_argument("--metadata", type=str, default="yolov8n_saved_model/metadata.yaml", help="Metadata yaml")
|
||||
args = parser.parse_args()
|
||||
|
||||
detector = YOLOv8TFLite(args.model, args.conf, args.iou, args.metadata)
|
||||
result = detector.detect(args.img)
|
||||
|
||||
cv2.imshow("Output", result)
|
||||
cv2.waitKey(0)
|
||||
123
algorithms/dms_yolo/code/examples/heatmaps.ipynb
Normal file
123
algorithms/dms_yolo/code/examples/heatmaps.ipynb
Normal file
@@ -0,0 +1,123 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PN1cAxdvd61e"
|
||||
},
|
||||
"source": [
|
||||
"<div align=\"center\">\n",
|
||||
"\n",
|
||||
" <a href=\"https://ultralytics.com/yolo\" target=\"_blank\">\n",
|
||||
" <img width=\"1024\", src=\"https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png\"></a>\n",
|
||||
"\n",
|
||||
" [中文](https://docs.ultralytics.com/zh/) | [한국어](https://docs.ultralytics.com/ko/) | [日本語](https://docs.ultralytics.com/ja/) | [Русский](https://docs.ultralytics.com/ru/) | [Deutsch](https://docs.ultralytics.com/de/) | [Français](https://docs.ultralytics.com/fr/) | [Español](https://docs.ultralytics.com/es/) | [Português](https://docs.ultralytics.com/pt/) | [Türkçe](https://docs.ultralytics.com/tr/) | [Tiếng Việt](https://docs.ultralytics.com/vi/) | [العربية](https://docs.ultralytics.com/ar/)\n",
|
||||
"\n",
|
||||
" <a href=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml\"><img src=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml/badge.svg\" alt=\"Ultralytics CI\"></a>\n",
|
||||
" <a href=\"https://console.paperspace.com/github/ultralytics/ultralytics\"><img src=\"https://assets.paperspace.io/img/gradient-badge.svg\" alt=\"Run on Gradient\"/></a>\n",
|
||||
" <a href=\"https://colab.research.google.com/github/ultralytics/ultralytics/blob/main/examples/heatmaps.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>\n",
|
||||
" <a href=\"https://www.kaggle.com/models/ultralytics/yolo26\"><img src=\"https://kaggle.com/static/images/open-in-kaggle.svg\" alt=\"Open In Kaggle\"></a>\n",
|
||||
" <a href=\"https://ultralytics.com/discord\"><img alt=\"Discord\" src=\"https://img.shields.io/discord/1089800235347353640?logo=discord&logoColor=white&label=Discord&color=blue\"></a>\n",
|
||||
"\n",
|
||||
"Welcome to the Ultralytics YOLO26 🚀 notebook! <a href=\"https://github.com/ultralytics/ultralytics\">YOLO26</a> is the latest version of the YOLO (You Only Look Once) AI models developed by <a href=\"https://ultralytics.com\">Ultralytics</a>. This notebook serves as the starting point for exploring the various resources available to help you get started with YOLO26 and understand its features and capabilities.\n",
|
||||
"\n",
|
||||
"YOLO26 models are fast, accurate, and easy to use, making them ideal for various object detection and image segmentation tasks. They can be trained on large datasets and run on diverse hardware platforms, from CPUs to GPUs.\n",
|
||||
"\n",
|
||||
"We hope that the resources in this notebook will help you get the most out of YOLO26. Please browse the YOLO26 <a href=\"https://docs.ultralytics.com/guides/heatmaps\">Heatmap Docs</a> for details, raise an issue on <a href=\"https://github.com/ultralytics/ultralytics\">GitHub</a> for support, and join our <a href=\"https://ultralytics.com/discord\">Discord</a> community for questions and discussions!\n",
|
||||
"\n",
|
||||
"</div>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "o68Sg1oOeZm2"
|
||||
},
|
||||
"source": [
|
||||
"# Setup\n",
|
||||
"\n",
|
||||
"pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/pyproject.toml) and check software and hardware.\n",
|
||||
"\n",
|
||||
"[](https://pypi.org/project/ultralytics/) [](https://clickpy.clickhouse.com/dashboard/ultralytics) [](https://pypi.org/project/ultralytics/)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9dSwz_uOReMI",
|
||||
"outputId": "99866c77-e210-41e1-d581-8508371ce634"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Ultralytics 8.3.174 🚀 Python-3.11.13 torch-2.6.0+cu124 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"Setup complete ✅ (2 CPUs, 12.7 GB RAM, 42.2/112.6 GB disk)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"!uv pip install ultralytics\n",
|
||||
"import ultralytics\n",
|
||||
"\n",
|
||||
"ultralytics.checks()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "m7VkxQ2aeg7k"
|
||||
},
|
||||
"source": [
|
||||
"# Introduction to Heatmaps\n",
|
||||
"\n",
|
||||
"A heatmap generated with [Ultralytics YOLO26](https://github.com/ultralytics/ultralytics/) transforms complex data into a vibrant, color-coded matrix. This visual tool employs a spectrum of colors to represent varying data values, where warmer hues indicate higher intensities and cooler tones signify lower values. Heatmaps excel in visualizing intricate data patterns, correlations, and anomalies, offering an accessible and engaging approach to data interpretation across diverse domains.\n",
|
||||
"\n",
|
||||
"## Real World Applications\n",
|
||||
"\n",
|
||||
"| Transportation | Retail |\n",
|
||||
"|:-----------------------------------------------------------------------------------------------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------:|\n",
|
||||
"|  |  |\n",
|
||||
"| Ultralytics YOLO11 Transportation Heatmap | Ultralytics YOLO11 Retail Heatmap |\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Cx-u59HQdu2o"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": "import cv2\n\nfrom ultralytics import solutions\n\n# Open video file\ncap = cv2.VideoCapture(\"path/to/video.mp4\")\nassert cap.isOpened(), \"Error reading video file\"\n\n# Get video properties\nw, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))\n\n# Initialize video writer\nvideo_writer = cv2.VideoWriter(\"heatmap_output.avi\", cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (w, h))\n\n# Initialize heatmap object\nheatmap_obj = solutions.Heatmap(\n colormap=cv2.COLORMAP_PARULA, # Color of the heatmap\n show=True, # Display the image during processing\n model=\"yolo26n.pt\", # Ultralytics YOLO26 model file\n)\n\nwhile cap.isOpened():\n success, im0 = cap.read()\n if not success:\n print(\"Video frame is empty or video processing has been successfully completed.\")\n break\n\n # Generate heatmap on the frame\n results = heatmap_obj(im0)\n\n # Write the frame to the output video\n video_writer.write(results.plot_im)\n\n# Release resources\ncap.release()\nvideo_writer.release()\ncv2.destroyAllWindows()"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QrlKg-y3fEyD"
|
||||
},
|
||||
"source": "# Additional Resources\n\n## Community Support\n\nFor more information on using heatmaps with Ultralytics, you can explore the comprehensive [Ultralytics Heatmaps Docs](https://docs.ultralytics.com/guides/heatmaps/). This guide covers everything from basic concepts to advanced techniques, ensuring you get the most out of your heatmap visualizations.\n\n## Ultralytics ⚡ Resources\n\nAt Ultralytics, we are committed to providing cutting-edge AI solutions. Here are some key resources to learn more about our company and get involved with our community:\n\n- [Ultralytics HUB](https://ultralytics.com/hub): Simplify your AI projects with Ultralytics HUB, our no-code tool for effortless YOLO training and deployment.\n- [Ultralytics Licensing](https://ultralytics.com/license): Review our licensing terms to understand how you can use our software in your projects.\n- [About Us](https://ultralytics.com/about): Discover our mission, vision, and the story behind Ultralytics.\n- [Join Our Team](https://ultralytics.com/work): Explore career opportunities and join our team of talented professionals.\n\n## YOLO26 🚀 Resources\n\nYOLO26 is the latest evolution in the YOLO series, offering state-of-the-art performance in object detection and image segmentation. Here are some essential resources to help you get started with YOLO26:\n\n- [GitHub](https://github.com/ultralytics/ultralytics): Access the YOLO26 repository on GitHub, where you can find the source code, contribute to the project, and report issues.\n- [Docs](https://docs.ultralytics.com/): Explore the official documentation for YOLO26, including installation guides, tutorials, and detailed API references.\n- [Discord](https://ultralytics.com/discord): Join our Discord community to connect with other users, share your projects, and get help from the Ultralytics team.\n\nThese resources are designed to help you leverage the full potential of Ultralytics' offerings and YOLO26. Whether you're a beginner or an experienced developer, you'll find the information and support you need to succeed."
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
115
algorithms/dms_yolo/code/examples/hub.ipynb
Normal file
115
algorithms/dms_yolo/code/examples/hub.ipynb
Normal file
@@ -0,0 +1,115 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "FIzICjaph_Wy"
|
||||
},
|
||||
"source": [
|
||||
"<a align=\"center\" href=\"https://ultralytics.com/hub\" target=\"_blank\">\n",
|
||||
"<img width=\"1024\", src=\"https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png\"></a>\n",
|
||||
"\n",
|
||||
"<div align=\"center\">\n",
|
||||
"\n",
|
||||
"[中文](https://docs.ultralytics.com/zh/hub/) | [한국어](https://docs.ultralytics.com/ko/hub/) | [日本語](https://docs.ultralytics.com/ja/hub/) | [Русский](https://docs.ultralytics.com/ru/hub/) | [Deutsch](https://docs.ultralytics.com/de/hub/) | [Français](https://docs.ultralytics.com/fr/hub/) | [Español](https://docs.ultralytics.com/es/hub/) | [Português](https://docs.ultralytics.com/pt/hub/) | [Türkçe](https://docs.ultralytics.com/tr/hub/) | [Tiếng Việt](https://docs.ultralytics.com/vi/hub/) | [العربية](https://docs.ultralytics.com/ar/hub/)\n",
|
||||
"\n",
|
||||
" <a href=\"https://github.com/ultralytics/hub/actions/workflows/ci.yml\"><img src=\"https://github.com/ultralytics/hub/actions/workflows/ci.yml/badge.svg\" alt=\"CI CPU\"></a>\n",
|
||||
" <a href=\"https://colab.research.google.com/github/ultralytics/hub/blob/main/hub.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>\n",
|
||||
"\n",
|
||||
" <a href=\"https://ultralytics.com/discord\"><img alt=\"Discord\" src=\"https://img.shields.io/discord/1089800235347353640?logo=discord&logoColor=white&label=Discord&color=blue\"></a>\n",
|
||||
" <a href=\"https://community.ultralytics.com\"><img alt=\"Ultralytics Forums\" src=\"https://img.shields.io/discourse/users?server=https%3A%2F%2Fcommunity.ultralytics.com&logo=discourse&label=Forums&color=blue\"></a>\n",
|
||||
" <a href=\"https://reddit.com/r/ultralytics\"><img alt=\"Ultralytics Reddit\" src=\"https://img.shields.io/reddit/subreddit-subscribers/ultralytics?style=flat&logo=reddit&logoColor=white&label=Reddit&color=blue\"></a>\n",
|
||||
"\n",
|
||||
"Welcome to the [Ultralytics](https://ultralytics.com/) HUB notebook!\n",
|
||||
"\n",
|
||||
"This notebook allows you to train Ultralytics [YOLO](https://github.com/ultralytics/ultralytics) 🚀 models using [HUB](https://hub.ultralytics.com/). Please browse the HUB <a href=\"https://docs.ultralytics.com/hub/\">Docs</a> for details, raise an issue on <a href=\"https://github.com/ultralytics/hub/issues/new/choose\">GitHub</a> for support, and join our <a href=\"https://ultralytics.com/discord\">Discord</a> community for questions and discussions!\n",
|
||||
"</div>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "eRQ2ow94MiOv"
|
||||
},
|
||||
"source": [
|
||||
"# Setup\n",
|
||||
"\n",
|
||||
"pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/pyproject.toml) and check software and hardware.\n",
|
||||
"\n",
|
||||
"[](https://pypi.org/project/ultralytics/) [](https://clickpy.clickhouse.com/dashboard/ultralytics) [](https://pypi.org/project/ultralytics/)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "FyDnXd-n4c7Y",
|
||||
"outputId": "e1d713ec-e8a6-4422-fe61-c76ec9f03df5"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Ultralytics 8.3.174 🚀 Python-3.11.13 torch-2.6.0+cu124 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"Setup complete ✅ (2 CPUs, 12.7 GB RAM, 42.2/112.6 GB disk)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"!uv pip install ultralytics # install\n",
|
||||
"from ultralytics import YOLO, checks, hub\n",
|
||||
"\n",
|
||||
"checks() # checks"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "cQ9BwaAqxAm4"
|
||||
},
|
||||
"source": [
|
||||
"# Start\n",
|
||||
"\n",
|
||||
"⚡ Login with your API key, load your YOLO 🚀 model and start training in 3 lines of code!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "XSlZaJ9Iw_iZ"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Log in to HUB using your API key (https://hub.ultralytics.com/settings?tab=api+keys)\n",
|
||||
"hub.login(\"YOUR_API_KEY\")\n",
|
||||
"\n",
|
||||
"# Load your model from HUB (replace 'YOUR_MODEL_ID' with your model ID)\n",
|
||||
"model = YOLO(\"https://hub.ultralytics.com/models/YOUR_MODEL_ID\")\n",
|
||||
"\n",
|
||||
"# Train the model\n",
|
||||
"results = model.train()"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"name": "Ultralytics HUB",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
130
algorithms/dms_yolo/code/examples/object_counting.ipynb
Normal file
130
algorithms/dms_yolo/code/examples/object_counting.ipynb
Normal file
@@ -0,0 +1,130 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PN1cAxdvd61e"
|
||||
},
|
||||
"source": [
|
||||
"<div align=\"center\">\n",
|
||||
"\n",
|
||||
" <a href=\"https://ultralytics.com/yolo\" target=\"_blank\">\n",
|
||||
" <img width=\"1024\", src=\"https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png\"></a>\n",
|
||||
"\n",
|
||||
" [中文](https://docs.ultralytics.com/zh/) | [한국어](https://docs.ultralytics.com/ko/) | [日本語](https://docs.ultralytics.com/ja/) | [Русский](https://docs.ultralytics.com/ru/) | [Deutsch](https://docs.ultralytics.com/de/) | [Français](https://docs.ultralytics.com/fr/) | [Español](https://docs.ultralytics.com/es/) | [Português](https://docs.ultralytics.com/pt/) | [Türkçe](https://docs.ultralytics.com/tr/) | [Tiếng Việt](https://docs.ultralytics.com/vi/) | [العربية](https://docs.ultralytics.com/ar/)\n",
|
||||
"\n",
|
||||
" <a href=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml\"><img src=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml/badge.svg\" alt=\"Ultralytics CI\"></a>\n",
|
||||
" <a href=\"https://console.paperspace.com/github/ultralytics/ultralytics\"><img src=\"https://assets.paperspace.io/img/gradient-badge.svg\" alt=\"Run on Gradient\"/></a>\n",
|
||||
" <a href=\"https://colab.research.google.com/github/ultralytics/ultralytics/blob/main/examples/object_counting.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>\n",
|
||||
" <a href=\"https://www.kaggle.com/models/ultralytics/yolo26\"><img src=\"https://kaggle.com/static/images/open-in-kaggle.svg\" alt=\"Open In Kaggle\"></a>\n",
|
||||
" <a href=\"https://ultralytics.com/discord\"><img alt=\"Discord\" src=\"https://img.shields.io/discord/1089800235347353640?logo=discord&logoColor=white&label=Discord&color=blue\"></a>\n",
|
||||
"\n",
|
||||
"Welcome to the Ultralytics YOLO26 🚀 notebook! <a href=\"https://github.com/ultralytics/ultralytics\">YOLO26</a> is the latest version of the YOLO (You Only Look Once) AI models developed by <a href=\"https://ultralytics.com\">Ultralytics</a>. This notebook serves as the starting point for exploring the various resources available to help you get started with YOLO26 and understand its features and capabilities.\n",
|
||||
"\n",
|
||||
"YOLO26 models are fast, accurate, and easy to use, making them ideal for various object detection and image segmentation tasks. They can be trained on large datasets and run on diverse hardware platforms, from CPUs to GPUs.\n",
|
||||
"\n",
|
||||
"We hope that the resources in this notebook will help you get the most out of YOLO26. Please browse the YOLO26 <a href=\"https://docs.ultralytics.com/guides/object-counting/\"> Object Counting Docs</a> for details, raise an issue on <a href=\"https://github.com/ultralytics/ultralytics\">GitHub</a> for support, and join our <a href=\"https://ultralytics.com/discord\">Discord</a> community for questions and discussions!\n",
|
||||
"\n",
|
||||
"</div>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "o68Sg1oOeZm2"
|
||||
},
|
||||
"source": [
|
||||
"# Setup\n",
|
||||
"\n",
|
||||
"pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/pyproject.toml) and check software and hardware.\n",
|
||||
"\n",
|
||||
"[](https://pypi.org/project/ultralytics/) [](https://clickpy.clickhouse.com/dashboard/ultralytics) [](https://pypi.org/project/ultralytics/)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9dSwz_uOReMI",
|
||||
"outputId": "fd3bab88-2f25-46c0-cae9-04d2beedc0c1"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Ultralytics 8.3.174 🚀 Python-3.11.13 torch-2.6.0+cu124 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"Setup complete ✅ (2 CPUs, 12.7 GB RAM, 42.2/112.6 GB disk)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"!uv pip install ultralytics\n",
|
||||
"import ultralytics\n",
|
||||
"\n",
|
||||
"ultralytics.checks()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "m7VkxQ2aeg7k"
|
||||
},
|
||||
"source": [
|
||||
"# Object Counting using Ultralytics YOLO26 🚀\n",
|
||||
"\n",
|
||||
"## What is Object Counting?\n",
|
||||
"\n",
|
||||
"Object counting with [Ultralytics YOLO11](https://github.com/ultralytics/ultralytics/) involves accurate identification and counting of specific objects in videos and camera streams. YOLO11 excels in real-time applications, providing efficient and precise object counting for various scenarios like crowd analysis and surveillance, thanks to its state-of-the-art algorithms and deep learning capabilities.\n",
|
||||
"\n",
|
||||
"## Advantages of Object Counting?\n",
|
||||
"\n",
|
||||
"- **Resource Optimization:** Object counting facilitates efficient resource management by providing accurate counts, and optimizing resource allocation in applications like inventory management.\n",
|
||||
"- **Enhanced Security:** Object counting enhances security and surveillance by accurately tracking and counting entities, aiding in proactive threat detection.\n",
|
||||
"- **Informed Decision-Making:** Object counting offers valuable insights for decision-making, optimizing processes in retail, traffic management, and various other domains.\n",
|
||||
"\n",
|
||||
"## Real World Applications\n",
|
||||
"\n",
|
||||
"| Logistics | Aquaculture |\n",
|
||||
"|:-------------------------------------------------------------------------------------------------------------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------:|\n",
|
||||
"|  |  |\n",
|
||||
"| Conveyor Belt Packets Counting Using Ultralytics YOLO26 | Fish Counting in Sea using Ultralytics YOLO26 |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Cx-u59HQdu2o"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": "import cv2\n\nfrom ultralytics import solutions\n\n# Open the video file\ncap = cv2.VideoCapture(\"path/to/video.mp4\")\nassert cap.isOpened(), \"Error reading video file\"\n\n# Get video properties: width, height, and frames per second (fps)\nw, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))\n\n# Define points for a line or region of interest in the video frame\nline_points = [(20, 400), (1080, 400)] # Line coordinates\n\n# Initialize the video writer to save the output video\nvideo_writer = cv2.VideoWriter(\"object_counting_output.avi\", cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (w, h))\n\n# Initialize the Object Counter with visualization options and other parameters\ncounter = solutions.ObjectCounter(\n show=True, # Display the image during processing\n region=line_points, # Region of interest points\n model=\"yolo26n.pt\", # Ultralytics YOLO26 model file\n line_width=2, # Thickness of the lines and bounding boxes\n)\n\n# Process video frames in a loop\nwhile cap.isOpened():\n success, im0 = cap.read()\n if not success:\n print(\"Video frame is empty or video processing has been successfully completed.\")\n break\n\n # Use the Object Counter to count objects in the frame and get the annotated image\n results = counter(im0)\n\n # Write the annotated frame to the output video\n video_writer.write(results.plot_im)\n\n# Release the video capture and writer objects\ncap.release()\nvideo_writer.release()\n\n# Close all OpenCV windows\ncv2.destroyAllWindows()"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QrlKg-y3fEyD"
|
||||
},
|
||||
"source": "# Additional Resources\n\n## Community Support\n\nFor more information on counting objects with Ultralytics, you can explore the comprehensive [Ultralytics Object Counting Docs](https://docs.ultralytics.com/guides/object-counting/). This guide covers everything from basic concepts to advanced techniques, ensuring you get the most out of counting and visualization.\n\n## Ultralytics ⚡ Resources\n\nAt Ultralytics, we are committed to providing cutting-edge AI solutions. Here are some key resources to learn more about our company and get involved with our community:\n\n- [Ultralytics HUB](https://ultralytics.com/hub): Simplify your AI projects with Ultralytics HUB, our no-code tool for effortless YOLO training and deployment.\n- [Ultralytics Licensing](https://ultralytics.com/license): Review our licensing terms to understand how you can use our software in your projects.\n- [About Us](https://ultralytics.com/about): Discover our mission, vision, and the story behind Ultralytics.\n- [Join Our Team](https://ultralytics.com/work): Explore career opportunities and join our team of talented professionals.\n\n## YOLO26 🚀 Resources\n\nYOLO26 is the latest evolution in the YOLO series, offering state-of-the-art performance in object detection and image segmentation. Here are some essential resources to help you get started with YOLO26:\n\n- [GitHub](https://github.com/ultralytics/ultralytics): Access the YOLO26 repository on GitHub, where you can find the source code, contribute to the project, and report issues.\n- [Docs](https://docs.ultralytics.com/): Explore the official documentation for YOLO26, including installation guides, tutorials, and detailed API references.\n- [Discord](https://ultralytics.com/discord): Join our Discord community to connect with other users, share your projects, and get help from the Ultralytics team.\n\nThese resources are designed to help you leverage the full potential of Ultralytics' offerings and YOLO26. Whether you're a beginner or an experienced developer, you'll find the information and support you need to succeed."
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
161
algorithms/dms_yolo/code/examples/object_tracking.ipynb
Normal file
161
algorithms/dms_yolo/code/examples/object_tracking.ipynb
Normal file
@@ -0,0 +1,161 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PN1cAxdvd61e"
|
||||
},
|
||||
"source": [
|
||||
"<div align=\"center\">\n",
|
||||
"\n",
|
||||
" <a href=\"https://ultralytics.com/yolo\" target=\"_blank\">\n",
|
||||
" <img width=\"1024\", src=\"https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png\"></a>\n",
|
||||
"\n",
|
||||
" [中文](https://docs.ultralytics.com/zh/) | [한국어](https://docs.ultralytics.com/ko/) | [日本語](https://docs.ultralytics.com/ja/) | [Русский](https://docs.ultralytics.com/ru/) | [Deutsch](https://docs.ultralytics.com/de/) | [Français](https://docs.ultralytics.com/fr/) | [Español](https://docs.ultralytics.com/es/) | [Português](https://docs.ultralytics.com/pt/) | [Türkçe](https://docs.ultralytics.com/tr/) | [Tiếng Việt](https://docs.ultralytics.com/vi/) | [العربية](https://docs.ultralytics.com/ar/)\n",
|
||||
"\n",
|
||||
" <a href=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml\"><img src=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml/badge.svg\" alt=\"Ultralytics CI\"></a>\n",
|
||||
" <a href=\"https://console.paperspace.com/github/ultralytics/ultralytics\"><img src=\"https://assets.paperspace.io/img/gradient-badge.svg\" alt=\"Run on Gradient\"/></a>\n",
|
||||
" <a href=\"https://colab.research.google.com/github/ultralytics/ultralytics/blob/main/examples/object_tracking.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>\n",
|
||||
" <a href=\"https://www.kaggle.com/models/ultralytics/yolo26\"><img src=\"https://kaggle.com/static/images/open-in-kaggle.svg\" alt=\"Open In Kaggle\"></a>\n",
|
||||
" <a href=\"https://ultralytics.com/discord\"><img alt=\"Discord\" src=\"https://img.shields.io/discord/1089800235347353640?logo=discord&logoColor=white&label=Discord&color=blue\"></a>\n",
|
||||
"\n",
|
||||
"Welcome to the Ultralytics YOLO26 🚀 notebook! <a href=\"https://github.com/ultralytics/ultralytics\">YOLO26</a> is the latest version of the YOLO (You Only Look Once) AI models developed by <a href=\"https://ultralytics.com\">Ultralytics</a>. This notebook serves as the starting point for exploring the various resources available to help you get started with YOLO26 and understand its features and capabilities.\n",
|
||||
"\n",
|
||||
"YOLO26 models are fast, accurate, and easy to use, making them ideal for various object detection and image segmentation tasks. They can be trained on large datasets and run on diverse hardware platforms, from CPUs to GPUs.\n",
|
||||
"\n",
|
||||
"We hope that the resources in this notebook will help you get the most out of YOLO26. Please browse the YOLO26 <a href=\"https://docs.ultralytics.com/modes/track/\"> Tracking Docs</a> for details, raise an issue on <a href=\"https://github.com/ultralytics/ultralytics\">GitHub</a> for support, and join our <a href=\"https://ultralytics.com/discord\">Discord</a> community for questions and discussions!\n",
|
||||
"\n",
|
||||
"</div>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "o68Sg1oOeZm2"
|
||||
},
|
||||
"source": [
|
||||
"# Setup\n",
|
||||
"\n",
|
||||
"pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/pyproject.toml) and check software and hardware.\n",
|
||||
"\n",
|
||||
"[](https://pypi.org/project/ultralytics/) [](https://clickpy.clickhouse.com/dashboard/ultralytics) [](https://pypi.org/project/ultralytics/)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9dSwz_uOReMI",
|
||||
"outputId": "ed8c2370-8fc7-4e4e-f669-d0bae4d944e9"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Ultralytics 8.3.174 🚀 Python-3.11.13 torch-2.6.0+cu124 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"Setup complete ✅ (2 CPUs, 12.7 GB RAM, 42.2/112.6 GB disk)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"!uv pip install ultralytics\n",
|
||||
"import ultralytics\n",
|
||||
"\n",
|
||||
"ultralytics.checks()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "m7VkxQ2aeg7k"
|
||||
},
|
||||
"source": [
|
||||
"# Ultralytics Object Tracking\n",
|
||||
"\n",
|
||||
"[Ultralytics YOLO11](https://github.com/ultralytics/ultralytics/) instance segmentation involves identifying and outlining individual objects in an image, providing a detailed understanding of spatial distribution. Unlike semantic segmentation, it uniquely labels and precisely delineates each object, crucial for tasks like object detection and medical imaging.\n",
|
||||
"\n",
|
||||
"There are two types of instance segmentation tracking available in the Ultralytics package:\n",
|
||||
"\n",
|
||||
"- **Instance Segmentation with Class Objects:** Each class object is assigned a unique color for clear visual separation.\n",
|
||||
"\n",
|
||||
"- **Instance Segmentation with Object Tracks:** Every track is represented by a distinct color, facilitating easy identification and tracking.\n",
|
||||
"\n",
|
||||
"## Samples\n",
|
||||
"\n",
|
||||
"| Instance Segmentation | Instance Segmentation + Object Tracking |\n",
|
||||
"|:---------------------------------------------------------------------------------------------------------------------------------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------:|\n",
|
||||
"|  |  |\n",
|
||||
"| Ultralytics Instance Segmentation 😍 | Ultralytics Instance Segmentation with Object Tracking 🔥 |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-ZF9DM6e6gz0"
|
||||
},
|
||||
"source": [
|
||||
"## CLI\n",
|
||||
"\n",
|
||||
"Command-Line Interface (CLI) example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-XJqhOwo6iqT"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!yolo track source=\"/path/to/video.mp4\" save=True"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "XRcw0vIE6oNb"
|
||||
},
|
||||
"source": [
|
||||
"## Python\n",
|
||||
"\n",
|
||||
"Python Instance Segmentation and Object tracking example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Cx-u59HQdu2o"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": "from collections import defaultdict\n\nimport cv2\n\nfrom ultralytics import YOLO\nfrom ultralytics.utils.plotting import Annotator, colors\n\n# Dictionary to store tracking history with default empty lists\ntrack_history = defaultdict(lambda: [])\n\n# Load the YOLO model with segmentation capabilities\nmodel = YOLO(\"yolo26n-seg.pt\")\n\n# Open the video file\ncap = cv2.VideoCapture(\"path/to/video.mp4\")\n\n# Retrieve video properties: width, height, and frames per second\nw, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))\n\n# Initialize video writer to save the output video with the specified properties\nout = cv2.VideoWriter(\"instance-segmentation-object-tracking.avi\", cv2.VideoWriter_fourcc(*\"MJPG\"), fps, (w, h))\n\nwhile True:\n # Read a frame from the video\n ret, im0 = cap.read()\n if not ret:\n print(\"Video frame is empty or video processing has been successfully completed.\")\n break\n\n # Create an annotator object to draw on the frame\n annotator = Annotator(im0, line_width=2)\n\n # Perform object tracking on the current frame\n results = model.track(im0, persist=True)\n\n # Check if tracking IDs and masks are present in the results\n if results[0].boxes.is_track and results[0].masks is not None:\n # Extract masks and tracking IDs\n masks = results[0].masks.xy\n track_ids = results[0].boxes.id.int().cpu().tolist()\n\n # Annotate each mask with its corresponding tracking ID and color\n for mask, track_id in zip(masks, track_ids):\n annotator.seg_bbox(mask=mask, mask_color=colors(int(track_id), True), label=str(track_id))\n\n # Write the annotated frame to the output video\n out.write(im0)\n # Display the annotated frame\n cv2.imshow(\"instance-segmentation-object-tracking\", im0)\n\n # Exit the loop if 'q' is pressed\n if cv2.waitKey(1) & 0xFF == ord(\"q\"):\n break\n\n# Release the video writer and capture objects, and close all OpenCV windows\nout.release()\ncap.release()\ncv2.destroyAllWindows()"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QrlKg-y3fEyD"
|
||||
},
|
||||
"source": "# Additional Resources\n\n## Community Support\n\nFor more information on using tracking with Ultralytics, you can explore the comprehensive [Ultralytics Tracking Docs](https://docs.ultralytics.com/modes/track/). This guide covers everything from basic concepts to advanced techniques, ensuring you get the most out of tracking and visualization.\n\n## Ultralytics ⚡ Resources\n\nAt Ultralytics, we are committed to providing cutting-edge AI solutions. Here are some key resources to learn more about our company and get involved with our community:\n\n- [Ultralytics HUB](https://ultralytics.com/hub): Simplify your AI projects with Ultralytics HUB, our no-code tool for effortless YOLO training and deployment.\n- [Ultralytics Licensing](https://ultralytics.com/license): Review our licensing terms to understand how you can use our software in your projects.\n- [About Us](https://ultralytics.com/about): Discover our mission, vision, and the story behind Ultralytics.\n- [Join Our Team](https://ultralytics.com/work): Explore career opportunities and join our team of talented professionals.\n\n## YOLO26 🚀 Resources\n\nYOLO26 is the latest evolution in the YOLO series, offering state-of-the-art performance in object detection and image segmentation. Here are some essential resources to help you get started with YOLO26:\n\n- [GitHub](https://github.com/ultralytics/ultralytics): Access the YOLO26 repository on GitHub, where you can find the source code, contribute to the project, and report issues.\n- [Docs](https://docs.ultralytics.com/): Explore the official documentation for YOLO26, including installation guides, tutorials, and detailed API references.\n- [Discord](https://ultralytics.com/discord): Join our Discord community to connect with other users, share your projects, and get help from the Ultralytics team.\n\nThese resources are designed to help you leverage the full potential of Ultralytics' offerings and YOLO26. Whether you're a beginner or an experienced developer, you'll find the information and support you need to succeed."
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
661
algorithms/dms_yolo/code/examples/tutorial.ipynb
Normal file
661
algorithms/dms_yolo/code/examples/tutorial.ipynb
Normal file
@@ -0,0 +1,661 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "YOLO26 Tutorial",
|
||||
"provenance": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "t6MPjfT5NrKQ"
|
||||
},
|
||||
"source": [
|
||||
"<div align=\"center\">\n",
|
||||
" <a href=\"https://ultralytics.com/yolo\" target=\"_blank\">\n",
|
||||
" <img width=\"1024\" src=\"https://raw.githubusercontent.com/ultralytics/assets/main/yolov8/banner-yolov8.png\">\n",
|
||||
" </a>\n",
|
||||
"\n",
|
||||
" [中文](https://docs.ultralytics.com/zh/) | [한국어](https://docs.ultralytics.com/ko/) | [日本語](https://docs.ultralytics.com/ja/) | [Русский](https://docs.ultralytics.com/ru/) | [Deutsch](https://docs.ultralytics.com/de/) | [Français](https://docs.ultralytics.com/fr/) | [Español](https://docs.ultralytics.com/es/) | [Português](https://docs.ultralytics.com/pt/) | [Türkçe](https://docs.ultralytics.com/tr/) | [Tiếng Việt](https://docs.ultralytics.com/vi/) | [العربية](https://docs.ultralytics.com/ar/)\n",
|
||||
"\n",
|
||||
" <a href=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml\"><img src=\"https://github.com/ultralytics/ultralytics/actions/workflows/ci.yml/badge.svg\" alt=\"Ultralytics CI\"></a>\n",
|
||||
" <a href=\"https://console.paperspace.com/github/ultralytics/ultralytics\"><img src=\"https://assets.paperspace.io/img/gradient-badge.svg\" alt=\"Run on Gradient\"/></a>\n",
|
||||
" <a href=\"https://colab.research.google.com/github/ultralytics/ultralytics/blob/main/examples/tutorial.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>\n",
|
||||
" <a href=\"https://www.kaggle.com/models/ultralytics/yolo26\"><img src=\"https://kaggle.com/static/images/open-in-kaggle.svg\" alt=\"Open In Kaggle\"></a>\n",
|
||||
"\n",
|
||||
" <a href=\"https://ultralytics.com/discord\"><img alt=\"Discord\" src=\"https://img.shields.io/discord/1089800235347353640?logo=discord&logoColor=white&label=Discord&color=blue\"></a>\n",
|
||||
" <a href=\"https://community.ultralytics.com\"><img alt=\"Ultralytics Forums\" src=\"https://img.shields.io/discourse/users?server=https%3A%2F%2Fcommunity.ultralytics.com&logo=discourse&label=Forums&color=blue\"></a>\n",
|
||||
" <a href=\"https://reddit.com/r/ultralytics\"><img alt=\"Ultralytics Reddit\" src=\"https://img.shields.io/reddit/subreddit-subscribers/ultralytics?style=flat&logo=reddit&logoColor=white&label=Reddit&color=blue\"></a>\n",
|
||||
"</div>\n",
|
||||
"\n",
|
||||
"This **Ultralytics Colab Notebook** is the easiest way to get started with [YOLO models](https://www.ultralytics.com/yolo)—no installation needed. Built by [Ultralytics](https://www.ultralytics.com/), the creators of YOLO, this notebook walks you through running **state-of-the-art** models directly in your browser.\n",
|
||||
"\n",
|
||||
"Ultralytics models are constantly updated for performance and flexibility. They're **fast**, **accurate**, and **easy to use**, and they excel at [object detection](https://docs.ultralytics.com/tasks/detect/), [tracking](https://docs.ultralytics.com/modes/track/), [instance segmentation](https://docs.ultralytics.com/tasks/segment/), [image classification](https://docs.ultralytics.com/tasks/classify/), and [pose estimation](https://docs.ultralytics.com/tasks/pose/).\n",
|
||||
"\n",
|
||||
"Find detailed documentation in the [Ultralytics Docs](https://docs.ultralytics.com/). Get support via [GitHub Issues](https://github.com/ultralytics/ultralytics/issues/new/choose). Join discussions on [Discord](https://discord.com/invite/ultralytics), [Reddit](https://www.reddit.com/r/ultralytics/), and the [Ultralytics Community Forums](https://community.ultralytics.com/)!\n",
|
||||
"\n",
|
||||
"Request an Enterprise License for commercial use at [Ultralytics Licensing](https://www.ultralytics.com/license)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7mGmQbAO5pQb"
|
||||
},
|
||||
"source": [
|
||||
"# Setup\n",
|
||||
"\n",
|
||||
"pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/pyproject.toml) and check software and hardware.\n",
|
||||
"\n",
|
||||
"[](https://pypi.org/project/ultralytics/) [](https://clickpy.clickhouse.com/dashboard/ultralytics) [](https://pypi.org/project/ultralytics/)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "wbvMlHd_QwMG",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "22044d50-807f-44dd-8768-98089231ee44"
|
||||
},
|
||||
"source": [
|
||||
"!uv pip install ultralytics\n",
|
||||
"import ultralytics\n",
|
||||
"ultralytics.checks()"
|
||||
],
|
||||
"execution_count": 1,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Ultralytics 8.4.4 🚀 Python-3.12.12 torch-2.9.0+cu126 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"Setup complete ✅ (2 CPUs, 12.7 GB RAM, 38.7/112.6 GB disk)\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4JnkELT0cIJg"
|
||||
},
|
||||
"source": [
|
||||
"# 1. Predict\n",
|
||||
"\n",
|
||||
"YOLO26 may be used directly in the Command Line Interface (CLI) with a `yolo` command for a variety of tasks and modes and accepts additional arguments, i.e. `imgsz=640`. See a full list of available `yolo` [arguments](https://docs.ultralytics.com/usage/cfg/) and other details in the [YOLO26 Predict Docs](https://docs.ultralytics.com/modes/predict/)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "zR9ZbuQCH7FX",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "7d35855e-42fa-4da4-a96e-0a5f09491b7d"
|
||||
},
|
||||
"source": [
|
||||
"# Run inference on an image with YOLO26n\n",
|
||||
"!yolo predict model=yolo26n.pt source='https://ultralytics.com/images/zidane.jpg'"
|
||||
],
|
||||
"execution_count": 2,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"\r\u001b[KDownloading https://github.com/ultralytics/assets/releases/download/v8.4.0/yolo26n.pt to 'yolo26n.pt': 100% ━━━━━━━━━━━━ 5.3MB 293.5MB/s 0.0s\n",
|
||||
"Ultralytics 8.4.4 🚀 Python-3.12.12 torch-2.9.0+cu126 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"YOLO26n summary (fused): 122 layers, 2,408,932 parameters, 0 gradients, 5.4 GFLOPs\n",
|
||||
"\n",
|
||||
"\u001b[KDownloading https://ultralytics.com/images/zidane.jpg to 'zidane.jpg': 100% ━━━━━━━━━━━━ 49.2KB 42.2MB/s 0.0s\n",
|
||||
"image 1/1 /content/zidane.jpg: 384x640 2 persons, 1 tie, 56.0ms\n",
|
||||
"Speed: 11.6ms preprocess, 56.0ms inference, 36.0ms postprocess per image at shape (1, 3, 384, 640)\n",
|
||||
"Results saved to \u001b[1m/content/runs/detect/predict\u001b[0m\n",
|
||||
"💡 Learn more at https://docs.ultralytics.com/modes/predict\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "hkAzDWJ7cWTr"
|
||||
},
|
||||
"source": [
|
||||
" \n",
|
||||
"<img align=\"left\" src=\"https://user-images.githubusercontent.com/26833433/212889447-69e5bdf1-5800-4e29-835e-2ed2336dede2.jpg\" width=\"600\">"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "0eq1SMWl6Sfn"
|
||||
},
|
||||
"source": [
|
||||
"# 2. Val\n",
|
||||
"Validate a model's accuracy on the [COCO](https://docs.ultralytics.com/datasets/detect/coco/) dataset's `val` or `test` splits. The latest YOLO26 [models](https://github.com/ultralytics/ultralytics#models) are downloaded automatically the first time they are used. See [YOLO26 Val Docs](https://docs.ultralytics.com/modes/val/) for more information."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "X58w8JLpMnjH"
|
||||
},
|
||||
"source": [
|
||||
"# Download COCO val\n",
|
||||
"from ultralytics.utils.downloads import download\n",
|
||||
"\n",
|
||||
"download('https://ultralytics.com/assets/coco2017val.zip', unzip=True, dir='datasets') # download (780MB - 5000 images)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZY2VXXXu74w5",
|
||||
"outputId": "1decfc78-03e0-4519-ac06-9596bd4dfcb9",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"# Validate YOLO26n on COCO8 val\n",
|
||||
"!yolo val model=yolo26n.pt data=coco8.yaml"
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Ultralytics 8.4.4 🚀 Python-3.12.12 torch-2.9.0+cu126 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"YOLO26n summary (fused): 122 layers, 2,408,932 parameters, 0 gradients, 5.4 GFLOPs\n",
|
||||
"\n",
|
||||
"WARNING ⚠️ Dataset 'coco8.yaml' images not found, missing path '/content/datasets/coco8/images/val'\n",
|
||||
"\u001b[KDownloading https://ultralytics.com/assets/coco8.zip to '/content/datasets/coco8.zip': 100% ━━━━━━━━━━━━ 432.8KB 36.5MB/s 0.0s\n",
|
||||
"\u001b[KUnzipping /content/datasets/coco8.zip to /content/datasets/coco8...: 100% ━━━━━━━━━━━━ 25/25 5.1Kfiles/s 0.0s\n",
|
||||
"Dataset download success ✅ (0.7s), saved to \u001b[1m/content/datasets\u001b[0m\n",
|
||||
"\n",
|
||||
"\u001b[KDownloading https://ultralytics.com/assets/Arial.ttf to '/root/.config/Ultralytics/Arial.ttf': 100% ━━━━━━━━━━━━ 755.1KB 106.1MB/s 0.0s\n",
|
||||
"\u001b[34m\u001b[1mval: \u001b[0mFast image access ✅ (ping: 0.0±0.0 ms, read: 1236.6±358.3 MB/s, size: 54.0 KB)\n",
|
||||
"\u001b[K\u001b[34m\u001b[1mval: \u001b[0mScanning /content/datasets/coco8/labels/val... 4 images, 0 backgrounds, 0 corrupt: 100% ━━━━━━━━━━━━ 4/4 103.0it/s 0.0s\n",
|
||||
"\u001b[34m\u001b[1mval: \u001b[0mNew cache created: /content/datasets/coco8/labels/val.cache\n",
|
||||
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 1/1 1.1s/it 1.1s\n",
|
||||
" all 4 17 0.849 0.654 0.907 0.667\n",
|
||||
" person 3 10 1 0.427 0.766 0.378\n",
|
||||
" dog 1 1 0.718 1 0.995 0.597\n",
|
||||
" horse 1 2 0.8 1 0.995 0.798\n",
|
||||
" elephant 1 2 0.854 0.5 0.695 0.438\n",
|
||||
" umbrella 1 1 0.724 1 0.995 0.895\n",
|
||||
" potted plant 1 1 1 0 0.995 0.895\n",
|
||||
"Speed: 1.0ms preprocess, 16.3ms inference, 0.0ms loss, 0.1ms postprocess per image\n",
|
||||
"Results saved to \u001b[1m/content/runs/detect/val\u001b[0m\n",
|
||||
"💡 Learn more at https://docs.ultralytics.com/modes/val\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"execution_count": 3
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# 3. Train\n",
|
||||
"\n",
|
||||
"<p align=\"\"><a href=\"https://ultralytics.com/hub\"><img width=\"1000\" src=\"https://github.com/ultralytics/assets/raw/main/yolov8/banner-integrations.png\"/></a></p>\n",
|
||||
"\n",
|
||||
"Train YOLO26 on [Detect](https://docs.ultralytics.com/tasks/detect/), [Segment](https://docs.ultralytics.com/tasks/segment/), [Classify](https://docs.ultralytics.com/tasks/classify/) and [Pose](https://docs.ultralytics.com/tasks/pose/) datasets. See [YOLO26 Train Docs](https://docs.ultralytics.com/modes/train/) for more information."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ktegpM42AooT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "1NcFxRcFdJ_O"
|
||||
},
|
||||
"source": [
|
||||
"#@title Select YOLO26 🚀 logger {run: 'auto'}\n",
|
||||
"logger = 'TensorBoard' #@param ['TensorBoard', 'Weights & Biases']\n",
|
||||
"\n",
|
||||
"if logger == 'TensorBoard':\n",
|
||||
" !yolo settings tensorboard=True\n",
|
||||
" %load_ext tensorboard\n",
|
||||
" %tensorboard --logdir .\n",
|
||||
"elif logger == 'Weights & Biases':\n",
|
||||
" !yolo settings wandb=True"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Train YOLO26n on COCO8 for 3 epochs\n",
|
||||
"!yolo train model=yolo26n.pt data=coco8.yaml epochs=3 imgsz=640"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "nPZZeNrLCQG6",
|
||||
"outputId": "b416bbb8-1249-4986-c9cf-0d1d04b59633",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Ultralytics 8.4.4 🚀 Python-3.12.12 torch-2.9.0+cu126 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"\u001b[34m\u001b[1mengine/trainer: \u001b[0magnostic_nms=False, amp=True, angle=1.0, augment=False, auto_augment=randaugment, batch=16, bgr=0.0, box=7.5, cache=False, cfg=None, classes=None, close_mosaic=10, cls=0.5, compile=False, conf=None, copy_paste=0.0, copy_paste_mode=flip, cos_lr=False, cutmix=0.0, data=coco8.yaml, degrees=0.0, deterministic=True, device=None, dfl=1.5, dnn=False, dropout=0.0, dynamic=False, embed=None, epochs=3, erasing=0.4, exist_ok=False, fliplr=0.5, flipud=0.0, format=torchscript, fraction=1.0, freeze=None, half=False, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, imgsz=640, int8=False, iou=0.7, keras=False, kobj=1.0, line_width=None, lr0=0.01, lrf=0.01, mask_ratio=4, max_det=300, mixup=0.0, mode=train, model=yolo26n.pt, momentum=0.937, mosaic=1.0, multi_scale=0.0, name=train, nbs=64, nms=False, opset=None, optimize=False, optimizer=auto, overlap_mask=True, patience=100, perspective=0.0, plots=True, pose=12.0, pretrained=True, profile=False, project=None, rect=False, resume=False, retina_masks=False, rle=1.0, save=True, save_conf=False, save_crop=False, save_dir=/content/runs/detect/train, save_frames=False, save_json=False, save_period=-1, save_txt=False, scale=0.5, seed=0, shear=0.0, show=False, show_boxes=True, show_conf=True, show_labels=True, simplify=True, single_cls=False, source=None, split=val, stream_buffer=False, task=detect, time=None, tracker=botsort.yaml, translate=0.1, val=True, verbose=True, vid_stride=1, visualize=False, warmup_bias_lr=0.1, warmup_epochs=3.0, warmup_momentum=0.8, weight_decay=0.0005, workers=8, workspace=None\n",
|
||||
"\n",
|
||||
" from n params module arguments \n",
|
||||
" 0 -1 1 464 ultralytics.nn.modules.conv.Conv [3, 16, 3, 2] \n",
|
||||
" 1 -1 1 4672 ultralytics.nn.modules.conv.Conv [16, 32, 3, 2] \n",
|
||||
" 2 -1 1 6640 ultralytics.nn.modules.block.C3k2 [32, 64, 1, False, 0.25] \n",
|
||||
" 3 -1 1 36992 ultralytics.nn.modules.conv.Conv [64, 64, 3, 2] \n",
|
||||
" 4 -1 1 26080 ultralytics.nn.modules.block.C3k2 [64, 128, 1, False, 0.25] \n",
|
||||
" 5 -1 1 147712 ultralytics.nn.modules.conv.Conv [128, 128, 3, 2] \n",
|
||||
" 6 -1 1 87040 ultralytics.nn.modules.block.C3k2 [128, 128, 1, True] \n",
|
||||
" 7 -1 1 295424 ultralytics.nn.modules.conv.Conv [128, 256, 3, 2] \n",
|
||||
" 8 -1 1 346112 ultralytics.nn.modules.block.C3k2 [256, 256, 1, True] \n",
|
||||
" 9 -1 1 164608 ultralytics.nn.modules.block.SPPF [256, 256, 5, 3, True] \n",
|
||||
" 10 -1 1 249728 ultralytics.nn.modules.block.C2PSA [256, 256, 1] \n",
|
||||
" 11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest'] \n",
|
||||
" 12 [-1, 6] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
||||
" 13 -1 1 119808 ultralytics.nn.modules.block.C3k2 [384, 128, 1, True] \n",
|
||||
" 14 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest'] \n",
|
||||
" 15 [-1, 4] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
||||
" 16 -1 1 34304 ultralytics.nn.modules.block.C3k2 [256, 64, 1, True] \n",
|
||||
" 17 -1 1 36992 ultralytics.nn.modules.conv.Conv [64, 64, 3, 2] \n",
|
||||
" 18 [-1, 13] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
||||
" 19 -1 1 95232 ultralytics.nn.modules.block.C3k2 [192, 128, 1, True] \n",
|
||||
" 20 -1 1 147712 ultralytics.nn.modules.conv.Conv [128, 128, 3, 2] \n",
|
||||
" 21 [-1, 10] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
||||
" 22 -1 1 463104 ultralytics.nn.modules.block.C3k2 [384, 256, 1, True, 0.5, True]\n",
|
||||
" 23 [16, 19, 22] 1 309656 ultralytics.nn.modules.head.Detect [80, 1, True, [64, 128, 256]] \n",
|
||||
"YOLO26n summary: 260 layers, 2,572,280 parameters, 2,572,280 gradients, 6.1 GFLOPs\n",
|
||||
"\n",
|
||||
"Transferred 708/708 items from pretrained weights\n",
|
||||
"\u001b[34m\u001b[1mAMP: \u001b[0mrunning Automatic Mixed Precision (AMP) checks...\n",
|
||||
"\u001b[34m\u001b[1mAMP: \u001b[0mchecks passed ✅\n",
|
||||
"\u001b[34m\u001b[1mtrain: \u001b[0mFast image access ✅ (ping: 0.0±0.0 ms, read: 1628.3±611.0 MB/s, size: 50.0 KB)\n",
|
||||
"\u001b[K\u001b[34m\u001b[1mtrain: \u001b[0mScanning /content/datasets/coco8/labels/train... 4 images, 0 backgrounds, 0 corrupt: 100% ━━━━━━━━━━━━ 4/4 1.8Kit/s 0.0s\n",
|
||||
"\u001b[34m\u001b[1mtrain: \u001b[0mNew cache created: /content/datasets/coco8/labels/train.cache\n",
|
||||
"\u001b[34m\u001b[1malbumentations: \u001b[0mBlur(p=0.01, blur_limit=(3, 7)), MedianBlur(p=0.01, blur_limit=(3, 7)), ToGray(p=0.01, method='weighted_average', num_output_channels=3), CLAHE(p=0.01, clip_limit=(1.0, 4.0), tile_grid_size=(8, 8))\n",
|
||||
"\u001b[34m\u001b[1mval: \u001b[0mFast image access ✅ (ping: 0.0±0.1 ms, read: 439.1±56.8 MB/s, size: 54.0 KB)\n",
|
||||
"\u001b[K\u001b[34m\u001b[1mval: \u001b[0mScanning /content/datasets/coco8/labels/val.cache... 4 images, 0 backgrounds, 0 corrupt: 100% ━━━━━━━━━━━━ 4/4 151.1Kit/s 0.0s\n",
|
||||
"Plotting labels to /content/runs/detect/train/labels.jpg... \n",
|
||||
"\u001b[34m\u001b[1moptimizer:\u001b[0m 'optimizer=auto' found, ignoring 'lr0=0.01' and 'momentum=0.937' and determining best 'optimizer', 'lr0' and 'momentum' automatically... \n",
|
||||
"\u001b[34m\u001b[1moptimizer:\u001b[0m MuSGD(lr=0.000119, momentum=0.9) with parameter groups 0 weight(decay=0.0), 0 weight(decay=0.0005), 24 bias(decay=0.0)\n",
|
||||
"Image sizes 640 train, 640 val\n",
|
||||
"Using 2 dataloader workers\n",
|
||||
"Logging results to \u001b[1m/content/runs/detect/train\u001b[0m\n",
|
||||
"Starting training for 3 epochs...\n",
|
||||
"\n",
|
||||
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
||||
"\u001b[K 1/3 0.75G 0.8467 1.76 0.01101 29 640: 100% ━━━━━━━━━━━━ 1/1 11.5s/it 11.5s\n",
|
||||
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 1/1 3.2s/it 3.2s\n",
|
||||
" all 4 17 0.9 0.653 0.901 0.668\n",
|
||||
"\n",
|
||||
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
||||
"\u001b[K 2/3 0.779G 1.332 2.559 0.01988 46 640: 100% ━━━━━━━━━━━━ 1/1 1.6it/s 0.6s\n",
|
||||
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 1/1 6.0it/s 0.2s\n",
|
||||
" all 4 17 0.895 0.655 0.902 0.667\n",
|
||||
"\n",
|
||||
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
||||
"\u001b[K 3/3 0.805G 0.9497 2.61 0.01205 20 640: 100% ━━━━━━━━━━━━ 1/1 3.7it/s 0.3s\n",
|
||||
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 1/1 7.4it/s 0.1s\n",
|
||||
" all 4 17 0.874 0.66 0.904 0.667\n",
|
||||
"\n",
|
||||
"3 epochs completed in 0.005 hours.\n",
|
||||
"Optimizer stripped from /content/runs/detect/train/weights/last.pt, 5.5MB\n",
|
||||
"Optimizer stripped from /content/runs/detect/train/weights/best.pt, 5.5MB\n",
|
||||
"\n",
|
||||
"Validating /content/runs/detect/train/weights/best.pt...\n",
|
||||
"Ultralytics 8.4.4 🚀 Python-3.12.12 torch-2.9.0+cu126 CUDA:0 (Tesla T4, 15095MiB)\n",
|
||||
"YOLO26n summary (fused): 122 layers, 2,408,932 parameters, 0 gradients, 5.4 GFLOPs\n",
|
||||
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 1/1 25.9it/s 0.0s\n",
|
||||
" all 4 17 0.9 0.653 0.901 0.668\n",
|
||||
" person 3 10 1 0.42 0.766 0.385\n",
|
||||
" dog 1 1 0.745 1 0.995 0.597\n",
|
||||
" horse 1 2 0.818 1 0.995 0.798\n",
|
||||
" elephant 1 2 0.84 0.5 0.662 0.434\n",
|
||||
" umbrella 1 1 1 1 0.995 0.895\n",
|
||||
" potted plant 1 1 1 0 0.995 0.895\n",
|
||||
"Speed: 0.3ms preprocess, 4.8ms inference, 0.0ms loss, 0.2ms postprocess per image\n",
|
||||
"Results saved to \u001b[1m/content/runs/detect/train\u001b[0m\n",
|
||||
"💡 Learn more at https://docs.ultralytics.com/modes/train\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"execution_count": 4
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# 4. Export\n",
|
||||
"\n",
|
||||
"Export a YOLO model to any supported format below with the `format` argument, i.e. `format=onnx`. See [Export Docs](https://docs.ultralytics.com/modes/export/) for more information.\n",
|
||||
"\n",
|
||||
"- 💡 ProTip: Export to [ONNX](https://docs.ultralytics.com/integrations/onnx/) or [OpenVINO](https://docs.ultralytics.com/integrations/openvino/) for up to 3x CPU speedup.\n",
|
||||
"- 💡 ProTip: Export to [TensorRT](https://docs.ultralytics.com/integrations/tensorrt/) for up to 5x GPU speedup.\n",
|
||||
"\n",
|
||||
"| Format | `format` Argument | Model | Metadata | Arguments |\n",
|
||||
"|--------|-----------------|-------|----------|------------|\n",
|
||||
"| [PyTorch](https://pytorch.org/) | - | `yolo26n.pt` | ✅ | - |\n",
|
||||
"| [TorchScript](https://docs.ultralytics.com/integrations/torchscript) | `torchscript` | `yolo26n.torchscript` | ✅ | `imgsz`, `batch`, `dynamic`, `optimize`, `half`, `nms`, `device` |\n",
|
||||
"| [ONNX](https://docs.ultralytics.com/integrations/onnx) | `onnx` | `yolo26n.onnx` | ✅ | `imgsz`, `batch`, `dynamic`, `half`, `opset`, `simplify`, `nms`, `device` |\n",
|
||||
"| [OpenVINO](https://docs.ultralytics.com/integrations/openvino) | `openvino` | `yolo26n_openvino_model/` | ✅ | `imgsz`, `batch`, `dynamic`, `half`, `int8`, `nms`, `fraction`, `device`, `data` |\n",
|
||||
"| [TensorRT](https://docs.ultralytics.com/integrations/tensorrt) | `engine` | `yolo26n.engine` | ✅ | `imgsz`, `batch`, `dynamic`, `half`, `int8`, `simplify`, `nms`, `fraction`, `device`, `data`, `workspace` |\n",
|
||||
"| [CoreML](https://docs.ultralytics.com/integrations/coreml) | `coreml` | `yolo26n.mlpackage` | ✅ | `imgsz`, `batch`, `half`, `int8`, `nms`, `device` |\n",
|
||||
"| [TF SavedModel](https://docs.ultralytics.com/integrations/tf-savedmodel) | `saved_model` | `yolo26n_saved_model/` | ✅ | `imgsz`, `batch`, `int8`, `keras`, `nms`, `device` |\n",
|
||||
"| [TF GraphDef](https://docs.ultralytics.com/integrations/tf-graphdef) | `pb` | `yolo26n.pb` | ❌ | `imgsz`, `batch`, `device` |\n",
|
||||
"| [TF Lite](https://docs.ultralytics.com/integrations/tflite) | `tflite` | `yolo26n.tflite` | ✅ | `imgsz`, `batch`, `half`, `int8`, `nms`, `fraction`, `device`, `data` |\n",
|
||||
"| [TF Edge TPU](https://docs.ultralytics.com/integrations/edge-tpu) | `edgetpu` | `yolo26n_edgetpu.tflite` | ✅ | `imgsz`, `device` |\n",
|
||||
"| [TF.js](https://docs.ultralytics.com/integrations/tfjs) | `tfjs` | `yolo26n_web_model/` | ✅ | `imgsz`, `batch`, `half`, `int8`, `nms`, `device` |\n",
|
||||
"| [PaddlePaddle](https://docs.ultralytics.com/integrations/paddlepaddle) | `paddle` | `yolo26n_paddle_model/` | ✅ | `imgsz`, `batch`, `device` |\n",
|
||||
"| [MNN](https://docs.ultralytics.com/integrations/mnn) | `mnn` | `yolo26n.mnn` | ✅ | `imgsz`, `batch`, `half`, `int8`, `device` |\n",
|
||||
"| [NCNN](https://docs.ultralytics.com/integrations/ncnn) | `ncnn` | `yolo26n_ncnn_model/` | ✅ | `imgsz`, `batch`, `half`, `device` |\n",
|
||||
"| [IMX500](https://docs.ultralytics.com/integrations/sony-imx500) | `imx` | `yolo26n_imx_model/` | ✅ | `imgsz`, `int8`, `fraction`, `device`, `data` |\n",
|
||||
"| [RKNN](https://docs.ultralytics.com/integrations/rockchip-rknn) | `rknn` | `yolo26n_rknn_model/` | ✅ | `imgsz`, `batch`, `name`, `device` |\n",
|
||||
"| [ExecuTorch](https://docs.ultralytics.com/integrations/executorch) | `executorch` | `yolo26n_executorch_model/` | ✅ | `imgsz`, `device` |\n",
|
||||
"| [Axelera](https://docs.ultralytics.com/integrations/axelera) | `axelera` | `yolo26n_axelera_model/` | ✅ | `imgsz`, `int8`, `fraction`, `device`, `data` |"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "CYIjW4igCjqD",
|
||||
"outputId": "9dd840c0-c689-49c6-f55a-b4cf366604fa"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!yolo export model=yolo26n.pt format=torchscript"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "kUMOQ0OeDBJG",
|
||||
"outputId": "57573520-046f-4479-a0d1-5953af7f3e9d",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Ultralytics 8.4.4 🚀 Python-3.12.12 torch-2.9.0+cu126 CPU (Intel Xeon CPU @ 2.00GHz)\n",
|
||||
"💡 ProTip: Export to OpenVINO format for best performance on Intel hardware. Learn more at https://docs.ultralytics.com/integrations/openvino/\n",
|
||||
"YOLO26n summary (fused): 122 layers, 2,408,932 parameters, 0 gradients, 5.4 GFLOPs\n",
|
||||
"\n",
|
||||
"\u001b[34m\u001b[1mPyTorch:\u001b[0m starting from 'yolo26n.pt' with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 300, 6) (5.3 MB)\n",
|
||||
"\n",
|
||||
"\u001b[34m\u001b[1mTorchScript:\u001b[0m starting export with torch 2.9.0+cu126...\n",
|
||||
"\u001b[34m\u001b[1mTorchScript:\u001b[0m export success ✅ 2.4s, saved as 'yolo26n.torchscript' (9.8 MB)\n",
|
||||
"\n",
|
||||
"Export complete (2.9s)\n",
|
||||
"Results saved to \u001b[1m/content\u001b[0m\n",
|
||||
"Predict: yolo predict task=detect model=yolo26n.torchscript imgsz=640 \n",
|
||||
"Validate: yolo val task=detect model=yolo26n.torchscript imgsz=640 data=/home/lq/codes/ultralytics/ultralytics/cfg/datasets/coco.yaml \n",
|
||||
"Visualize: https://netron.app\n",
|
||||
"💡 Learn more at https://docs.ultralytics.com/modes/export\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"execution_count": 5
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# 5. Python Usage\n",
|
||||
"\n",
|
||||
"YOLO26 was reimagined using Python-first principles for the most seamless Python YOLO experience yet. YOLO26 models can be loaded from a trained checkpoint or created from scratch. Then methods are used to train, val, predict, and export the model. See detailed Python usage examples in the [YOLO26 Python Docs](https://docs.ultralytics.com/usage/python/)."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "bpF9-vS_DAaf"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from ultralytics import YOLO\n",
|
||||
"\n",
|
||||
"# Load a model\n",
|
||||
"model = YOLO('yolo26n.yaml') # build a new model from scratch\n",
|
||||
"model = YOLO('yolo26n.pt') # load a pretrained model (recommended for training)\n",
|
||||
"\n",
|
||||
"# Use the model\n",
|
||||
"results = model.train(data='coco8.yaml', epochs=3) # train the model\n",
|
||||
"results = model.val() # evaluate model performance on the validation set\n",
|
||||
"results = model('https://ultralytics.com/images/bus.jpg') # predict on an image\n",
|
||||
"results = model.export(format='onnx') # export the model to ONNX format"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Phm9ccmOKye5"
|
||||
},
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# 6. Tasks\n",
|
||||
"\n",
|
||||
"YOLO26 can train, val, predict and export models for the most common tasks in vision AI: [Detect](https://docs.ultralytics.com/tasks/detect/), [Segment](https://docs.ultralytics.com/tasks/segment/), [Classify](https://docs.ultralytics.com/tasks/classify/) and [Pose](https://docs.ultralytics.com/tasks/pose/). See [YOLO26 Tasks Docs](https://docs.ultralytics.com/tasks/) for more information.\n",
|
||||
"\n",
|
||||
"<br><img width=\"1024\" src=\"https://raw.githubusercontent.com/ultralytics/assets/main/im/banner-tasks.png\">"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "yq26lwpYK1lq"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## 1. Detection\n",
|
||||
"\n",
|
||||
"YOLO26 _detection_ models have no suffix and are the default YOLO26 models, i.e. `yolo26n.pt` and are pretrained on COCO. See [Detection Docs](https://docs.ultralytics.com/tasks/detect/) for full details."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "8Go5qqS9LbC5"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Load YOLO26n, train it on COCO128 for 3 epochs and predict an image with it\n",
|
||||
"from ultralytics import YOLO\n",
|
||||
"\n",
|
||||
"model = YOLO('yolo26n.pt') # load a pretrained YOLO detection model\n",
|
||||
"model.train(data='coco8.yaml', epochs=3) # train the model\n",
|
||||
"model('https://ultralytics.com/images/bus.jpg') # predict on an image"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "7ZW58jUzK66B"
|
||||
},
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## 2. Segmentation\n",
|
||||
"\n",
|
||||
"YOLO26 _segmentation_ models use the `-seg` suffix, i.e. `yolo26n-seg.pt` and are pretrained on COCO. See [Segmentation Docs](https://docs.ultralytics.com/tasks/segment/) for full details."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "WFPJIQl_L5HT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Load YOLO26n-seg, train it on COCO128-seg for 3 epochs and predict an image with it\n",
|
||||
"from ultralytics import YOLO\n",
|
||||
"\n",
|
||||
"model = YOLO('yolo26n-seg.pt') # load a pretrained YOLO segmentation model\n",
|
||||
"model.train(data='coco8-seg.yaml', epochs=3) # train the model\n",
|
||||
"model('https://ultralytics.com/images/bus.jpg') # predict on an image"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ax3p94VNK9zR"
|
||||
},
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## 3. Classification\n",
|
||||
"\n",
|
||||
"YOLO26 _classification_ models use the `-cls` suffix, i.e. `yolo26n-cls.pt` and are pretrained on ImageNet. See [Classification Docs](https://docs.ultralytics.com/tasks/classify/) for full details."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "5q9Zu6zlL5rS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Load YOLO26n-cls, train it on mnist160 for 3 epochs and predict an image with it\n",
|
||||
"from ultralytics import YOLO\n",
|
||||
"\n",
|
||||
"model = YOLO('yolo26n-cls.pt') # load a pretrained YOLO classification model\n",
|
||||
"model.train(data='mnist160', epochs=3) # train the model\n",
|
||||
"model('https://ultralytics.com/images/bus.jpg') # predict on an image"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SpIaFLiO11TG"
|
||||
},
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## 4. Pose\n",
|
||||
"\n",
|
||||
"YOLO26 _pose_ models use the `-pose` suffix, i.e. `yolo26n-pose.pt` and are pretrained on COCO Keypoints. See [Pose Docs](https://docs.ultralytics.com/tasks/pose/) for full details."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "si4aKFNg19vX"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Load YOLO26n-pose, train it on COCO8-pose for 3 epochs and predict an image with it\n",
|
||||
"from ultralytics import YOLO\n",
|
||||
"\n",
|
||||
"model = YOLO('yolo26n-pose.pt') # load a pretrained YOLO pose model\n",
|
||||
"model.train(data='coco8-pose.yaml', epochs=3) # train the model\n",
|
||||
"model('https://ultralytics.com/images/bus.jpg') # predict on an image"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "cf5j_T9-B5F0"
|
||||
},
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## 5. Oriented Bounding Boxes (OBB)\n",
|
||||
"\n",
|
||||
"YOLO26 _OBB_ models use the `-obb` suffix, i.e. `yolo26n-obb.pt` and are pretrained on the DOTA dataset. See [OBB Docs](https://docs.ultralytics.com/tasks/obb/) for full details."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "IJNKClOOB5YS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "IEijrePND_2I"
|
||||
},
|
||||
"source": [
|
||||
"# Load YOLO26n-obb, train it on DOTA8 for 3 epochs and predict an image with it\n",
|
||||
"from ultralytics import YOLO\n",
|
||||
"\n",
|
||||
"model = YOLO('yolo26n-obb.pt') # load a pretrained YOLO OBB model\n",
|
||||
"model.train(data='dota8.yaml', epochs=3) # train the model\n",
|
||||
"model('https://ultralytics.com/images/boats.jpg') # predict on an image"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Appendix\n",
|
||||
"\n",
|
||||
"Additional content below."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "pIdE6i8C3LYp"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Pip install from source\n",
|
||||
"!uv pip install git+https://github.com/ultralytics/ultralytics@main"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "uRKlwxSJdhd1"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Git clone and run tests on 'main' branch\n",
|
||||
"!git clone https://github.com/ultralytics/ultralytics -b main\n",
|
||||
"!uv pip install -qe ultralytics"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "GtPlh7mcCGZX"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Run tests (Git clone only)\n",
|
||||
"!pytest ultralytics/tests"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Wdc6t_bfzDDk"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"metadata": {
|
||||
"id": "_icJkqGnXFv8"
|
||||
},
|
||||
"cell_type": "code",
|
||||
"outputs": [],
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"# Validate multiple models\n",
|
||||
"for x in 'nsmlx':\n",
|
||||
" !yolo val model=yolo26{x}.pt data=coco.yaml"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
Reference in New Issue
Block a user