1个不稳定版本
0.1.0 | 2023年11月17日 |
---|
#2101 在 算法
46KB
420 行
适用于OpenCV生态系统中基于YOLO神经网络的Rust编程语言的对象检测工具
此crate提供了一些基本结构和方法,用于通过OpenCV的DNN模块解决对象检测任务。目前实施和测试的工作流程
网络类型 | Darknet | ONNX |
---|---|---|
YOLO v3 tiny | ✅ | ❌ |
YOLO v4 tiny | ✅ | ❌ |
YOLO v7 tiny | ✅ | ❌ |
YOLO v3 | ❓ (需要测试) | ❌ |
YOLO v4 | ❓ (需要测试) | ❌ |
YOLO v7 | ❓ (需要测试) | ❌ |
YOLO v8 n | ❌ (是否可行?) | ✅ |
目录
关于
- 为什么?
嗯,我在我的私人项目和公共项目中都对模板代码(模型初始化、后处理函数等)感到有点厌倦。
- 何时有用?
嗯,有几种情况你可能需要这个crate
- 你需要使用YOLO作为你的神经网络基础;
- 你不想使用Pytorch / Tensorflow / Jax或任何其他DL/ML框架(将来可能会在这个crate中使用纯ONNX,而无需OpenCV功能 - PR欢迎);
- 你需要使用OpenCV的DNN模块来初始化神经网络;
- 为什么没有YOLOv5?
我认为v8和v5版本的后处理有一些区别。我需要更多时间来调查v5应该如何工作。
- 测试了哪个OpenCV版本?
我使用v4.7.0进行了测试。Rust绑定版本:v0.66.0
- 包装结构是否线程安全?
我不确定它是否打算在多个线程中使用(PR欢迎)。但如果你想给提供的结构提供“异步”访问,我认为你应该使用某种队列机制。
先决条件
-
当然,你必须安装OpenCV并带有DNN扩展模块。如果你需要利用GPU/OpenVINO的强大功能,那么你也应该考虑相应的扩展模块。
我喜欢根据我的需求对它进行一些调整(OpenCV版本/启用Python绑定)来使用这个Makefile。
-
准备神经网络:训练它或获取预训练模型。我提供了简单的Bash脚本来下载“小型”版本的YOLO
- YOLO v3 tiny - download_v3_tiny.sh;
- YOLO v4 tiny - download_v4_tiny.sh;
- YOLO v7 tiny - download_v7_tiny.sh;
- YOLO v8 nano (n) - download_v8_n.sh. 注意,此脚本下载基于Pytorch的权重,并通过Python的
包将其转换为ONNX格式。
使用方法
有一些示例,但我将一步一步引导你
-
将此crate添加到您的
Cargo.toml
cargo add od_opencv
-
也将OpenCV的绑定crate添加到
Cargo.toml
# I'm using 0.66 version cargo add [email protected]
-
下载预训练模型或使用您自己的神经网络。
我将使用先决条件部分中的预训练权重
-
在您的
main.rs
文件中导入“基本”OpenCV库use opencv::{ core::{Scalar, Vector}, imgcodecs::imread, imgcodecs::imwrite, imgproc::LINE_4, imgproc::rectangle, dnn::DNN_BACKEND_CUDA, // I will utilize my GPU to perform faster inference. Your way may vary dnn::DNN_TARGET_CUDA, };
-
导入crate
use od_opencv::{ model_format::ModelFormat, // I'll use YOLOv8 by Ultralytics. // If you prefer traditional YOLO, then import it as: // model_classic::ModelYOLOClassic model_ultralytics::ModelUltralyticsV8 };
-
准备模型
// Define classes (in this case we consider 80 COCO labels) let classes_labels: Vec<&str> = vec!["person", "bicycle", "car", "motorbike", "aeroplane", "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", "sofa", "pottedplant", "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"]; // Define format for OpenCV's DNN module let mf = ModelFormat::ONNX; // Define model's input size let net_width = 640; let net_height = 640; // Initialize model itself let mut model = ModelUltralyticsV8::new_from_file("pretrained/yolov8n.onnx", None, (net_width, net_height), mf, DNN_BACKEND_CUDA, DNN_TARGET_CUDA, vec![]).unwrap(); // Read image into the OpenCV's Mat object // Define it as mutable since we are going to put bounding boxes onto it. let mut frame = imread("images/dog.jpg", 1).unwrap(); // Feed forward image through the model let (bboxes, class_ids, confidences) = model.forward(&frame, 0.25, 0.4).unwrap(); // Process results for (i, bbox) in bboxes.iter().enumerate() { // Place bounding boxes onto the image rectangle(&mut frame, *bbox, Scalar::from((0.0, 255.0, 0.0)), 2, LINE_4, 0).unwrap(); // Debug output to stdin println!("Class: {}", classes_labels[class_ids[i]]); println!("\tBounding box: {:?}", bbox); println!("\tConfidences: {}", confidences[i]); } // Finally save the updated image to the file system imwrite("images/dog_yolov8_n.jpg", &frame, &Vector::new()).unwrap();
-
您已经准备好开始了
cargo run
-
如果遇到任何问题,请随时创建一个issue
参考
- YOLO v3论文 - https://arxiv.org/abs/1804.02767,Joseph Redmon,Ali Farhadi
- YOLO v4论文 - https://arxiv.org/abs/2004.10934,Alexey Bochkovskiy,Chien-Yao Wang,Hong-Yuan Mark Liao
- YOLO v7论文 - https://arxiv.org/abs/2207.02696,Chien-Yao Wang,Alexey Bochkovskiy,Hong-Yuan Mark Liao
- 原始Darknet YOLO仓库 - https://github.com/pjreddie/darknet
- Darknet YOLO最受欢迎的分支 - https://github.com/AlexeyAB/darknet
- YOLOv8的开发者 - https://github.com/ultralytics/ultralytics。如果您知道关于YOLOv8架构的一些原始论文,请与我联系,以便在README中提及。
- Rust OpenCV的绑定 - https://github.com/twistedfall/opencv-rust
- Go OpenCV的绑定(用于现成的Makefile) - https://github.com/hybridgroup/gocv
依赖关系
~1.7–3MB
~25K SLoC