yolo_binding\core/
mod.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
mod export;
mod load;
mod predict;
use std::{collections::HashMap, error::Error};
use tch::{CModule, Tensor};
pub enum YOLODevice {
    Cpu,
    Gpu,
}
pub struct YOLO {
    yolo_model: CModule,
    cuda: YOLODevice,
    pub types: HashMap<i64, String>,
}

impl YOLO {
    pub fn new(model_path: &str, cuda_is_available: bool) -> Self {
        //! Create a new YOLO model
        //! model_path: the path of the model
        //! cuda_is_available: whether to use cuda
        load::load_model_from_path(model_path, cuda_is_available).unwrap()
    }

    pub fn predict(&self, input: &Tensor) -> Result<Tensor, Box<dyn Error>> {
        //! Predict the input tensor
        //! input: the input tensor
        //! return: the output tensor
        predict::pred(self, input)
    }
}
pub fn load_one_image(image_path: &str) -> Result<Tensor, Box<dyn Error>> {
    //! Load one image from the path
    //! image_path: the path of the image
    //! return: the image tensor
    let image = load::load_one_image(image_path)?;
    Ok(image)
}
pub fn load_images_from_dir(image_dir: &str) -> Result<Tensor, Box<dyn Error>> {
    //! Load images from the directory
    //! image_dir: the directory of the images
    //! return: the images tensor
    let images = load::load_images_from_dir(image_dir)?;
    Ok(images)
}
pub fn load_one_image_from_memory(image_bytes: &[u8]) -> Result<Tensor, Box<dyn Error>> {
    //! Load one image from the memory
    //! image_bytes: the bytes of the image
    //! return: the image tensor
    let image = load::load_one_image_from_memory(image_bytes)?;
    Ok(image)
}
pub fn get_results(
    input: &Tensor,
    confidence: f64,
    threshold: f64,
) -> Result<Vec<Vec<(i64, i64, i64, i64, i64, f64)>>, Box<dyn Error>> {
    //! Get the results from the output tensor
    //! input: the output tensor from the model [n, X, 8400]
    //! confidence: the confidence threshold
    //! threshold: the NMS threshold
    //! return: the results of the bounding boxes
    //! x, y, w, h, class, confidence
    let mut result = Vec::new();
    let (batch_size, _, _) = input.size3().unwrap();
    for i in 0..batch_size {
        let res = export::post_process(&input.get(i), confidence, threshold)?;
        result.push(res);
    }
    Ok(result)
}