Rust编程语言中针对OpenCV生态系统中YOLO神经网络对象检测的实用工具

此crate提供了一些基本的结构和方法，用于通过OpenCV的DNN模块解决对象检测任务。目前实现和测试的工作流程

目录

• 关于
• 先决条件
• 使用方法
• 参考

关于

- 为什么？

嗯，我在我的私人和公开项目中都尝试过一些样板代码（模型初始化、后处理函数等）。

- 何时有用？

嗯，在某些情况下，您可能需要这个crate

• 您需要将YOLO作为您的神经网络基础；
• 您不想使用Pytorch / Tensorflow / Jax或任何其他DL/ML框架（将来可能会在这个crate中使用纯ONNX，而不使用OpenCV功能 - 欢迎PR）；
• 您需要使用OpenCV的DNN模块来初始化神经网络；

- 为什么没有YOLOv5？

我认为v8和v5版本之间的后处理功能有所不同。我需要更多时间来调查如何确切地使其工作。

- 测试了哪个OpenCV版本？

我已经与v4.7.0一起测试了它。Rust绑定版本：v0.66.0

- 包装结构是否线程安全？

我不确定它是否打算在多个线程中使用（欢迎PR）。但如果你想给提供的结构提供“异步”访问，我认为你应该使用一些队列机制。

先决条件

• 当然，您必须安装带有DNN扩展模块的OpenCV。如果您需要利用GPU/OpenVINO的功能，那么您还需要考虑包含相应的扩展模块。

  我喜欢根据我的需求对它进行微调（OpenCV的版本/启用Python绑定）来使用这个Makefile。

• 准备神经网络：训练它或获取预训练的模型。我提供了简单的Bash脚本来下载"小型"版本的YOLO。

  • YOLO v3 tiny - download_v3_tiny.sh；YOLO v3 - download_v3.sh；
  • YOLO v4 tiny - download_v4_tiny.sh；YOLO v4 - download_v4.sh；
  • YOLO v7 tiny - download_v7_tiny.sh；YOLO v7 - download_v7.sh；
  • YOLO v8 nano (n) - download_v8_n.sh。
  • YOLO v8 small (s) - download_v8_s.sh。
  • YOLO v8 medium (m) - download_v8_m.sh。
  • YOLO v8 large (l) - download_v8_l.sh。
  • YOLO v8 extra (x) - download_v8_x.sh。

  注意："v8"脚本下载基于Pytorch的权重，并通过Python的包将其转换为ONNX权重。

使用方法

有一些示例，但我会一步步引导你。

1. 将此crate添加到你的Cargo.toml中。

   cargo add od_opencv
   

2. 也将OpenCV的绑定crate添加到Cargo.toml中。

   # I'm using 0.66 version
   cargo add opencv@0.66
   

3. 下载预训练模型或使用你自己的神经网络。

   我将使用先决条件部分中的预训练权重。

4. 在您的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,
   };
   

5. 导入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
   };
   

6. 准备模型。

   // 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 optional filters.
   // E.g.: if you do want to find only dogs and cats and you can't re-train neural network, 
   // then you can just place `vec![15, 16]` to filter dogs and cats (15 - index of `cat` in class labels, 16 - `dog`)
   // let class_filters: Vec<usize> = vec![15, 16];
   let class_filters: Vec<usize> = vec![];
   
   // 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, class_filters).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();
   

7. 你现在可以开始了。

   cargo run
   

8. 如果出现问题，请随时创建一个问题。

参考

• 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