使用 Rust 和 OpenCV 進行物體檢測

ttocr、com發表於2024-12-03

我們將利用 Rust 的 opencv 庫來封裝 OpenCV 的功能,載入 YOLOv3 模型並進行影像的物體檢測。YOLO(You Only Look Once)是一個非常高效的實時目標檢測模型,能夠在圖片中快速定位出物體並標註出邊界框。

環境準備

  1. 安裝 Rust
    如果你還沒有安裝 Rust,可以使用下面的命令進行安裝:

bash

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
然後透過以下命令更新 Rust:

bash

rustup update
2. 安裝 OpenCV
為了在 Rust 中使用 OpenCV,我們需要安裝 OpenCV 庫,並透過 opencv crate 來訪問它。首先,確保你已經安裝了 OpenCV(可以參考 OpenCV 官方安裝文件 進行安裝)。

在 Ubuntu 上,可以使用以下命令安裝 OpenCV:

bash

sudo apt update
sudo apt install libopencv-dev
3. 設定 Cargo.toml
在專案的 Cargo.toml 檔案中,新增 opencv 依賴項:

toml

[dependencies]
opencv = "0.67"
然後執行以下命令來下載依賴:

bash

cargo build
4. 下載 YOLO 模型和配置檔案
下載 YOLOv3 配置檔案:yolov3.cfg
下載 YOLOv3 權重檔案:yolov3.weights
程式碼實現
下面是一個用 Rust 和 OpenCV 實現 YOLOv3 物體檢測的示例。

rust

extern crate opencv;

use opencv::{
core::{Mat, Scalar, Size},
dnn::{read_net, Net},
highgui::{imshow, wait_key, named_window},
imgcodecs::imread,
imgproc::{cvt_color, COLOR_BGR2RGB},
prelude:😗,
types::{VectorOfMat, VectorOfString},
};

const CONFIDENCE_THRESHOLD: f32 = 0.5;
const NMS_THRESHOLD: f32 = 0.4;
const INPUT_WIDTH: i32 = 416;
const INPUT_HEIGHT: i32 = 416;

fn main() -> opencv::Result<()> {
// 載入 YOLO 模型
let model_path = "yolov3.weights";
let config_path = "yolov3.cfg";
let net = read_net(config_path, model_path)?;

// 載入類別標籤
let class_names = vec![
    "person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat", "traffic light",
];

// 載入輸入影像
let image = imread("input_image.jpg", opencv::imgcodecs::IMREAD_COLOR)?;
let mut blob = Mat::default();
cvt_color(&image, &mut blob, COLOR_BGR2RGB, 0)?;
let blob = dnn::blob_from_image(&image, 1.0 / 255.0, Size::new(INPUT_WIDTH, INPUT_HEIGHT), Scalar::new(0.0, 0.0, 0.0, 0.0), true, false)?;

// 將輸入資料傳遞給網路
net.set_input(&blob, "", 1.0, Scalar::new(0.0, 0.0, 0.0, 0.0));

// 獲取網路輸出層
let output_layer_names = net.get_unconnected_out_layers_names()?;

// 前向推理
let mut outs = VectorOfMat::new();
net.forward(&mut outs, &output_layer_names)?;

// 後處理,提取物體框
let (boxes, confidences, class_ids) = post_process(&image, &outs);

// 對物體框進行非最大抑制
let indices = dnn::nms_boxes(&boxes, &confidences, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)?;

// 繪製物體框
for i in 0..indices.len() {
    let box_ = boxes[indices[i] as usize];
    let label = format!(
        "{}: {:.2}%",
        class_names[class_ids[indices[i] as usize] as usize],
        confidences[indices[i] as usize] * 100.0
    );

    opencv::imgproc::rectangle(
        &mut image,
        box_,
        Scalar::new(0.0, 255.0, 0.0, 0.0),
        2,
        8,
        0,
    )?;

    opencv::imgproc::put_text(
        &mut image,
        &label,
        opencv::core::Point::new(box_.x, box_.y - 10),
        opencv::imgproc::FONT_HERSHEY_SIMPLEX,
        0.5,
        Scalar::new(0.0, 255.0, 0.0, 0.0),
        1,
        8,
        false,
    )?;
}

// 顯示結果
named_window("Object Detection", opencv::highgui::WINDOW_NORMAL)?;
imshow("Object Detection", &image)?;
wait_key(0)?;

Ok(())

}

// 後處理:提取物體框、置信度和類別
fn post_process(image: &Mat, outs: &VectorOfMat) -> (Veccv::core::Rect, Vec, Vec) {
let mut boxes = Vec::new();
let mut confidences = Vec::new();
let mut class_ids = Vec::new();

for i in 0..outs.len() {
    let mat = outs.get(i)?.clone();
    let mut data = mat.data()?;
    let rows = mat.rows();
    let cols = mat.cols();

    for r in 0..rows {
        let score = &data[(r * cols)..];
        let confidence = score[5];

        if confidence > CONFIDENCE_THRESHOLD {
            let center_x = (score[0] * image.cols() as f32) as i32;
            let center_y = (score[1] * image.rows() as f32) as i32;
            let width = (score[2] * image.cols() as f32) as i32;
            let height = (score[3] * image.rows() as f32) as i32;
            let box_ = cv::core::Rect::new(center_x - width / 2, center_y - height / 2, width, height);

            boxes.push(box_);
            confidences.push(confidence);
            class_ids.push(score[4] as i32);
        }
    }
}更多內容訪問ttocr.com或聯絡1436423940

(boxes, confidences, class_ids)

}
程式碼解析
載入 YOLO 模型和配置檔案: 使用 opencv::dnn::read_net() 函式載入 YOLO 配置檔案和權重檔案,形成一個深度神經網路物件。

影像預處理: 透過 opencv::dnn::blob_from_image() 將輸入影像轉換為 YOLO 模型需要的格式。

前向推理: 使用 net.forward() 對輸入影像進行前向推理,獲取物體檢測的結果。

後處理: 提取預測結果,包括物體框、置信度和類別,並進行非最大抑制(NMS)。

繪製結果: 使用 OpenCV 繪製檢測到的物體框,並標註出類別和置信度。

相關文章