├── .gitignore ├── .dockerignore ├── data └── config.json ├── Cargo.toml ├── LICENSE ├── README.md ├── Dockerfile ├── src ├── main.rs └── model.rs └── Cargo.lock /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | *.jpg -------------------------------------------------------------------------------- /.dockerignore: -------------------------------------------------------------------------------- 1 | target/ 2 | data/ 3 | *.jpg 4 | *.lock -------------------------------------------------------------------------------- /data/config.json: -------------------------------------------------------------------------------- 1 | { 2 | 3 | "model_path": "your-model.onnx", 4 | "input_size": 1280, 5 | "class_names": ["Class 1", "Class 2", "Class 3"] 6 | 7 | } -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "yolov5-api" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | 10 | serde = { version = "1.0", features = ["derive"] } 11 | serde_json = "1.0" 12 | opencv = "0.66" 13 | actix-web = "4" -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Mauro Sciancalepore 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # YoloV5-API [WIP] 2 | 3 | API to run inferences with YoloV5 models. Written in Rust, based on OpenCV 4.5.5 4 | 5 | If you need a C++ version, check my [C++ Yolov5-API](https://github.com/masc-it/yolov5-api-cpp) 6 | 7 | ## Requirements 8 | 9 | - [OpenCV 4.5.5](https://github.com/opencv/opencv/releases/tag/4.5.5) installed on your system 10 | - Follow [Rust opencv README](https://github.com/twistedfall/opencv-rust) 11 | 12 | ## Model config 13 | 14 | **Data** directory must contain your config.json 15 | 16 | **config.json** defines: 17 | - ONNX absolute model path 18 | - input size (640 default) 19 | - array of class names 20 | 21 | A dummy example is available in the _data/_ folder 22 | 23 | 24 | ## Docker 25 | 26 | docker run --name rust-yolov5-api -v :/app/data -p 5000:5000 mascit/rust-yolov5-api:latest 27 | 28 | ## Build 29 | 30 | Development: 31 | 32 | cargo build 33 | 34 | Release: 35 | 36 | cargo build --release 37 | 38 | ## Run 39 | 40 | cargo run 41 | 42 | For Windows users: Assure to have _opencv_world455.dll_ in your exe directory. 43 | 44 | # Endpoints 45 | 46 | ## /predict [POST] 47 | 48 | ### Body 49 | - Image bytes (binary in Postman) 50 | 51 | ### Headers 52 | - X-Confidence-Thresh 53 | - default 0.5 54 | - X-NMS-Thresh 55 | - default 0.45 56 | - X-Return 57 | - image_with_boxes 58 | - A JPEG image with drawn predictions 59 | - json (default) 60 | - A json array containing predictions. Each object defines: xmin, ymin, xmax, ymax, conf, class_name -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM debian:stable-slim 2 | ENV TZ=Europe/London 3 | RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone 4 | 5 | RUN apt update -y && apt-get install -y build-essential \ 6 | pkg-config \ 7 | gcc \ 8 | cmake \ 9 | libavcodec-dev \ 10 | libavformat-dev \ 11 | libswscale-dev \ 12 | libv4l-dev \ 13 | libxvidcore-dev \ 14 | libx264-dev \ 15 | libjpeg-dev \ 16 | libpng-dev \ 17 | libtiff-dev \ 18 | libatlas-base-dev \ 19 | libtbb2 \ 20 | libtbb-dev \ 21 | wget \ 22 | curl \ 23 | clang \ 24 | unzip 25 | 26 | RUN wget -O opencv.zip "https://github.com/opencv/opencv/archive/4.5.5.zip" \ 27 | && unzip "opencv.zip" \ 28 | && mkdir -p build && cd build 29 | 30 | RUN cmake -DBUILD_opencv_java=OFF \ 31 | -DWITH_QT=OFF -DWITH_GTK=OFF \ 32 | -DBUILD_opencv_python=OFF \ 33 | -DWITH_CUDA=OFF \ 34 | -DWITH_OPENGL=OFF \ 35 | -DWITH_OPENCL=ON \ 36 | -DWITH_IPP=ON \ 37 | -DWITH_TBB=ON \ 38 | -DWITH_EIGEN=ON \ 39 | -DWITH_V4L=ON \ 40 | -DBUILD_TESTS=OFF \ 41 | -DBUILD_PERF_TESTS=OFF \ 42 | -DOPENCV_GENERATE_PKGCONFIG=YES \ 43 | -DCMAKE_BUILD_TYPE=RELEASE \ 44 | -DCMAKE_INSTALL_PREFIX=/usr/local \ 45 | "../opencv-4.5.5" 46 | 47 | #RUN cmake --build . -j 10 48 | RUN make install -j 10 49 | 50 | RUN pkg-config --cflags opencv4 && pkg-config --libs opencv4 51 | 52 | ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/lib" 53 | WORKDIR /app 54 | RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y 55 | ENV PATH="/root/.cargo/bin:${PATH}" 56 | 57 | RUN rm -rf ../opencv-4.5.5 && rm ../opencv.zip 58 | 59 | COPY . . 60 | 61 | RUN cargo build --release 62 | 63 | #EXPOSE 5000 64 | CMD ["./target/release/yolov5-api"] -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | 2 | use std::time::Instant; 3 | use std::error::Error; 4 | use std::sync::Mutex; 5 | 6 | use model::Model; 7 | use actix_web::{get, post, web, App, HttpResponse, HttpServer, Responder, HttpRequest}; 8 | 9 | mod model; 10 | 11 | #[get("/")] 12 | async fn hello() -> impl Responder { 13 | HttpResponse::Ok().body("Hello world!") 14 | } 15 | 16 | #[post("/predict")] 17 | async fn predict(data: web::Data, req: HttpRequest, req_body: web::Bytes) -> impl Responder { 18 | 19 | let start = Instant::now(); 20 | let mut model = &mut *data.model.lock().unwrap(); 21 | 22 | let img: Vec = req_body.to_vec(); 23 | 24 | let conf_thresh = req.headers().get("X-Confidence-Thresh") 25 | .map_or_else(|| 0.5, 26 | |f| f.to_str().unwrap().parse::().map_or(0.5, |v| v)); 27 | 28 | 29 | let nms_thresh = req.headers().get("X-NMS-Thresh") 30 | .map_or_else(|| 0.5, 31 | |f| f.to_str().unwrap().parse::().map_or(0.5, |v| v)); 32 | 33 | let mut return_type = req.headers().get("X-Return") 34 | .map_or_else(|| "json", |f| f.to_str().unwrap()); 35 | 36 | if !return_type.eq("json") && !return_type.eq("img_with_boxes") { 37 | return_type = "json"; 38 | } 39 | let img_vec: opencv::core::Vector = opencv::core::Vector::from_iter(img); 40 | let mat = opencv::imgcodecs::imdecode(&img_vec, opencv::imgcodecs::IMREAD_UNCHANGED); 41 | 42 | let mut mat = mat.unwrap(); 43 | 44 | let detections = model::detect(&mut model,&mat, conf_thresh, nms_thresh); 45 | 46 | if detections.is_err() { 47 | return HttpResponse::Ok() 48 | .append_header(("Content-Type", "application/json")) 49 | .body("{\"msg\": \"Invalid image.\"}") 50 | } 51 | 52 | let detections = detections.unwrap(); 53 | let duration = start.elapsed(); 54 | if return_type.eq("json") { 55 | let json_response = serde_json::to_string_pretty(&detections).unwrap(); 56 | HttpResponse::Ok() 57 | .append_header(("X-Duration", duration.as_millis().to_string())) 58 | .append_header(("Content-Type", "application/json")) 59 | .body(json_response) 60 | } else { 61 | 62 | let img_with_boxes_bytes = model::draw_predictions(&mut mat, &detections, &model.model_config).unwrap(); 63 | 64 | HttpResponse::Ok() 65 | .append_header(("X-Duration", duration.as_millis().to_string())) 66 | .append_header(("Content-Type", "image/jpeg")) 67 | .body(img_with_boxes_bytes) 68 | } 69 | 70 | 71 | 72 | 73 | } 74 | 75 | struct AppState { 76 | 77 | model: Mutex 78 | } 79 | 80 | #[actix_web::main] 81 | async fn main_api(model: Model) -> std::io::Result<()> { 82 | 83 | let data = web::Data::new(AppState { 84 | model: Mutex::new(model), 85 | }); 86 | 87 | HttpServer::new(move || { 88 | App::new() 89 | .app_data(data.clone()) 90 | .service(hello) 91 | .service(predict) 92 | }) 93 | .bind(("0.0.0.0", 5000))? 94 | .run() 95 | .await 96 | } 97 | 98 | fn main() -> Result<(), Box>{ 99 | let model = model::load_model()?; 100 | 101 | main_api(model).unwrap(); 102 | 103 | Ok(()) 104 | } 105 | -------------------------------------------------------------------------------- /Cargo.lock: -------------------------------------------------------------------------------- 1 | # This file is automatically @generated by Cargo. 2 | # It is not intended for manual editing. 3 | version = 3 4 | 5 | [[package]] 6 | name = "aho-corasick" 7 | version = "0.7.18" 8 | source = "registry+https://github.com/rust-lang/crates.io-index" 9 | checksum = "1e37cfd5e7657ada45f742d6e99ca5788580b5c529dc78faf11ece6dc702656f" 10 | dependencies = [ 11 | "memchr", 12 | ] 13 | 14 | [[package]] 15 | name = "autocfg" 16 | version = "1.1.0" 17 | source = "registry+https://github.com/rust-lang/crates.io-index" 18 | checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" 19 | 20 | [[package]] 21 | name = "cc" 22 | version = "1.0.73" 23 | source = "registry+https://github.com/rust-lang/crates.io-index" 24 | checksum = "2fff2a6927b3bb87f9595d67196a70493f627687a71d87a0d692242c33f58c11" 25 | dependencies = [ 26 | "jobserver", 27 | ] 28 | 29 | [[package]] 30 | name = "cfg-if" 31 | version = "1.0.0" 32 | source = "registry+https://github.com/rust-lang/crates.io-index" 33 | checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" 34 | 35 | [[package]] 36 | name = "clang" 37 | version = "2.0.0" 38 | source = "registry+https://github.com/rust-lang/crates.io-index" 39 | checksum = "84c044c781163c001b913cd018fc95a628c50d0d2dfea8bca77dad71edb16e37" 40 | dependencies = [ 41 | "clang-sys", 42 | "libc", 43 | ] 44 | 45 | [[package]] 46 | name = "clang-sys" 47 | version = "1.3.3" 48 | source = "registry+https://github.com/rust-lang/crates.io-index" 49 | checksum = "5a050e2153c5be08febd6734e29298e844fdb0fa21aeddd63b4eb7baa106c69b" 50 | dependencies = [ 51 | "glob", 52 | "libc", 53 | "libloading", 54 | ] 55 | 56 | [[package]] 57 | name = "dunce" 58 | version = "1.0.2" 59 | source = "registry+https://github.com/rust-lang/crates.io-index" 60 | checksum = "453440c271cf5577fd2a40e4942540cb7d0d2f85e27c8d07dd0023c925a67541" 61 | 62 | [[package]] 63 | name = "glob" 64 | version = "0.3.0" 65 | source = "registry+https://github.com/rust-lang/crates.io-index" 66 | checksum = "9b919933a397b79c37e33b77bb2aa3dc8eb6e165ad809e58ff75bc7db2e34574" 67 | 68 | [[package]] 69 | name = "jobserver" 70 | version = "0.1.24" 71 | source = "registry+https://github.com/rust-lang/crates.io-index" 72 | checksum = "af25a77299a7f711a01975c35a6a424eb6862092cc2d6c72c4ed6cbc56dfc1fa" 73 | dependencies = [ 74 | "libc", 75 | ] 76 | 77 | [[package]] 78 | name = "libc" 79 | version = "0.2.126" 80 | source = "registry+https://github.com/rust-lang/crates.io-index" 81 | checksum = "349d5a591cd28b49e1d1037471617a32ddcda5731b99419008085f72d5a53836" 82 | 83 | [[package]] 84 | name = "libloading" 85 | version = "0.7.3" 86 | source = "registry+https://github.com/rust-lang/crates.io-index" 87 | checksum = "efbc0f03f9a775e9f6aed295c6a1ba2253c5757a9e03d55c6caa46a681abcddd" 88 | dependencies = [ 89 | "cfg-if", 90 | "winapi", 91 | ] 92 | 93 | [[package]] 94 | name = "maplit" 95 | version = "1.0.2" 96 | source = "registry+https://github.com/rust-lang/crates.io-index" 97 | checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d" 98 | 99 | [[package]] 100 | name = "memchr" 101 | version = "2.5.0" 102 | source = "registry+https://github.com/rust-lang/crates.io-index" 103 | checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" 104 | 105 | [[package]] 106 | name = "num-traits" 107 | version = "0.2.15" 108 | source = "registry+https://github.com/rust-lang/crates.io-index" 109 | checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd" 110 | dependencies = [ 111 | "autocfg", 112 | ] 113 | 114 | [[package]] 115 | name = "once_cell" 116 | version = "1.13.0" 117 | source = "registry+https://github.com/rust-lang/crates.io-index" 118 | checksum = "18a6dbe30758c9f83eb00cbea4ac95966305f5a7772f3f42ebfc7fc7eddbd8e1" 119 | 120 | [[package]] 121 | name = "opencv" 122 | version = "0.66.0" 123 | source = "registry+https://github.com/rust-lang/crates.io-index" 124 | checksum = "ae1211402fae55d0054e5779bcb231660752b8da825f88aa8a837aef7944c68d" 125 | dependencies = [ 126 | "cc", 127 | "clang", 128 | "dunce", 129 | "glob", 130 | "jobserver", 131 | "libc", 132 | "num-traits", 133 | "once_cell", 134 | "opencv-binding-generator", 135 | "pkg-config", 136 | "semver", 137 | "shlex", 138 | "vcpkg", 139 | ] 140 | 141 | [[package]] 142 | name = "opencv-binding-generator" 143 | version = "0.45.0" 144 | source = "registry+https://github.com/rust-lang/crates.io-index" 145 | checksum = "1070da825ea3584b7ef0a2951fc002843f986ec430681b2273d9d425da2ffd61" 146 | dependencies = [ 147 | "clang", 148 | "clang-sys", 149 | "dunce", 150 | "maplit", 151 | "once_cell", 152 | "percent-encoding", 153 | "regex", 154 | ] 155 | 156 | [[package]] 157 | name = "percent-encoding" 158 | version = "2.1.0" 159 | source = "registry+https://github.com/rust-lang/crates.io-index" 160 | checksum = "d4fd5641d01c8f18a23da7b6fe29298ff4b55afcccdf78973b24cf3175fee32e" 161 | 162 | [[package]] 163 | name = "pkg-config" 164 | version = "0.3.25" 165 | source = "registry+https://github.com/rust-lang/crates.io-index" 166 | checksum = "1df8c4ec4b0627e53bdf214615ad287367e482558cf84b109250b37464dc03ae" 167 | 168 | [[package]] 169 | name = "regex" 170 | version = "1.6.0" 171 | source = "registry+https://github.com/rust-lang/crates.io-index" 172 | checksum = "4c4eb3267174b8c6c2f654116623910a0fef09c4753f8dd83db29c48a0df988b" 173 | dependencies = [ 174 | "aho-corasick", 175 | "memchr", 176 | "regex-syntax", 177 | ] 178 | 179 | [[package]] 180 | name = "regex-syntax" 181 | version = "0.6.27" 182 | source = "registry+https://github.com/rust-lang/crates.io-index" 183 | checksum = "a3f87b73ce11b1619a3c6332f45341e0047173771e8b8b73f87bfeefb7b56244" 184 | 185 | [[package]] 186 | name = "semver" 187 | version = "1.0.12" 188 | source = "registry+https://github.com/rust-lang/crates.io-index" 189 | checksum = "a2333e6df6d6598f2b1974829f853c2b4c5f4a6e503c10af918081aa6f8564e1" 190 | 191 | [[package]] 192 | name = "shlex" 193 | version = "1.1.0" 194 | source = "registry+https://github.com/rust-lang/crates.io-index" 195 | checksum = "43b2853a4d09f215c24cc5489c992ce46052d359b5109343cbafbf26bc62f8a3" 196 | 197 | [[package]] 198 | name = "vcpkg" 199 | version = "0.2.15" 200 | source = "registry+https://github.com/rust-lang/crates.io-index" 201 | checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" 202 | 203 | [[package]] 204 | name = "winapi" 205 | version = "0.3.9" 206 | source = "registry+https://github.com/rust-lang/crates.io-index" 207 | checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" 208 | dependencies = [ 209 | "winapi-i686-pc-windows-gnu", 210 | "winapi-x86_64-pc-windows-gnu", 211 | ] 212 | 213 | [[package]] 214 | name = "winapi-i686-pc-windows-gnu" 215 | version = "0.4.0" 216 | source = "registry+https://github.com/rust-lang/crates.io-index" 217 | checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" 218 | 219 | [[package]] 220 | name = "winapi-x86_64-pc-windows-gnu" 221 | version = "0.4.0" 222 | source = "registry+https://github.com/rust-lang/crates.io-index" 223 | checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" 224 | 225 | [[package]] 226 | name = "yolov5-api" 227 | version = "0.1.0" 228 | dependencies = [ 229 | "opencv", 230 | ] 231 | -------------------------------------------------------------------------------- /src/model.rs: -------------------------------------------------------------------------------- 1 | use opencv::{ 2 | core::{self, MatTraitConst, MatTrait, MatExprTraitConst}, 3 | dnn::{self, NetTraitConst, NetTrait} 4 | }; 5 | use std::{fs::File, io::{BufReader}, error::Error}; 6 | use serde::{Deserialize, Serialize}; 7 | use std::os::raw::c_void; 8 | 9 | #[derive(Serialize, Deserialize)] 10 | pub struct BoxDetection { 11 | 12 | pub xmin: i32, 13 | pub ymin: i32, 14 | pub xmax: i32, 15 | pub ymax: i32, 16 | 17 | pub class: i32, 18 | pub conf: f32 19 | 20 | } 21 | 22 | #[derive(Serialize, Deserialize)] 23 | pub struct Detections { 24 | 25 | pub detections: Vec 26 | } 27 | 28 | /// Defines model information. 29 | /// - ONNX model absolute path 30 | /// - array of class names 31 | /// - model input size 32 | #[derive(Deserialize)] 33 | pub struct ModelConfig { 34 | 35 | pub model_path : String, 36 | pub class_names : Vec, 37 | pub input_size: i32 38 | } 39 | 40 | /// Contains the instantiated model itself and its configuration. 41 | pub struct Model { 42 | 43 | pub model: dnn::Net, 44 | pub model_config: ModelConfig 45 | } 46 | 47 | /// Contains information about original input image and effective size fed into the model. 48 | pub struct MatInfo { 49 | 50 | width: f32, 51 | height: f32, 52 | scaled_size: f32 53 | } 54 | 55 | /// Run detection on input image. 56 | pub fn detect(model_data: &mut Model, img: &core::Mat, conf_thresh: f32, nms_thresh: f32) -> opencv::Result { 57 | 58 | let model = &mut model_data.model; 59 | 60 | let model_config = &mut model_data.model_config; 61 | 62 | let mat_info = MatInfo{ 63 | width: img.cols() as f32, 64 | height: img.rows() as f32, 65 | 66 | scaled_size: model_config.input_size as f32 67 | }; 68 | 69 | let padded_mat = prepare_input(&img).unwrap(); 70 | 71 | // dnn blob 72 | let blob = opencv::dnn::blob_from_image(&padded_mat, 1.0 / 255.0, opencv::core::Size_{width: model_config.input_size, height: model_config.input_size}, core::Scalar::new(0f64,0f64,0f64,0f64), true, false, core::CV_32F)?; 73 | 74 | let out_layer_names = model.get_unconnected_out_layers_names()?; 75 | 76 | let mut outs : opencv::core::Vector = opencv::core::Vector::default(); 77 | model.set_input(&blob, "", 1.0, core::Scalar::default())?; 78 | 79 | model.forward(&mut outs, &out_layer_names)?; 80 | 81 | let detections = post_process(&outs, &mat_info, conf_thresh, nms_thresh)?; 82 | 83 | Ok(detections) 84 | } 85 | 86 | /// Prepare an image as a squared matrix. 87 | fn prepare_input(img: &core::Mat) -> opencv::Result { 88 | 89 | let width = img.cols(); 90 | let height = img.rows(); 91 | 92 | let _max = std::cmp::max(width, height); 93 | 94 | let mut result = opencv::core::Mat::zeros(_max, _max, opencv::core::CV_8UC3).unwrap().to_mat().unwrap(); 95 | 96 | img.copy_to(&mut result)?; 97 | 98 | Ok(result) 99 | 100 | } 101 | 102 | /// Process predictions and apply NMS. 103 | fn post_process(outs: &core::Vector, mat_info: &MatInfo, conf_thresh: f32, nms_thresh: f32 ) -> opencv::Result{ 104 | 105 | let mut det = outs.get(0)?; 106 | 107 | let rows = *det.mat_size().get(1).unwrap(); 108 | let cols = *det.mat_size().get(2).unwrap(); 109 | 110 | let mut boxes: core::Vector = core::Vector::default(); 111 | let mut scores: core::Vector = core::Vector::default(); 112 | 113 | let mut indices: core::Vector = core::Vector::default(); 114 | 115 | let mut class_index_list: core::Vector = core::Vector::default(); 116 | 117 | let x_factor = mat_info.width / mat_info.scaled_size; 118 | let y_factor = mat_info.height / mat_info.scaled_size; 119 | 120 | unsafe { 121 | 122 | let mut min_val = Some(0f64); 123 | let mut max_val = Some(0f64); 124 | 125 | let mut min_loc = Some(core::Point::default()); 126 | let mut max_loc = Some(core::Point::default()); 127 | let mut mask = core::no_array(); 128 | 129 | let data = det.ptr_mut(0)?.cast::(); 130 | 131 | // safe alternative needed.. 132 | let m = core::Mat::new_rows_cols_with_data(rows, cols, core::CV_32F, data, core::Mat_AUTO_STEP )?; 133 | 134 | for r in 0..m.rows() { 135 | 136 | let cx: &f32 = m.at_2d::(r, 0)?; 137 | let cy: &f32 = m.at_2d::(r, 1)?; 138 | let w: &f32 = m.at_2d::(r, 2)?; 139 | let h: &f32 = m.at_2d::(r, 3)?; 140 | let sc: &f32 = m.at_2d::(r, 4)?; 141 | 142 | let score = *sc as f64; 143 | 144 | if score < conf_thresh.into() { 145 | continue; 146 | } 147 | let confs = m.row(r)?.col_range( &core::Range::new(5, m.row(r)?.cols())?)?; 148 | 149 | let c = (confs * score).into_result()?.to_mat()?; 150 | 151 | // find predicted class with highest confidence 152 | core::min_max_loc(&c, min_val.as_mut(), max_val.as_mut(), min_loc.as_mut(), max_loc.as_mut(), &mut mask)?; 153 | 154 | scores.push(max_val.unwrap() as f32); 155 | class_index_list.push(max_loc.unwrap().x); 156 | 157 | boxes.push( core::Rect{ 158 | x: (((*cx) - (*w) / 2.0) * x_factor).round() as i32, 159 | y: (((*cy) - (*h) / 2.0) * y_factor).round() as i32, 160 | width: (*w * x_factor).round() as i32, 161 | height: (*h * y_factor).round() as i32 162 | } ); 163 | indices.push(r); 164 | 165 | } 166 | 167 | } 168 | 169 | // Run NMS. 170 | dnn::nms_boxes(&boxes, &scores, conf_thresh, nms_thresh, &mut indices, 1.0, 0)?; 171 | 172 | let mut final_boxes : Vec = Vec::default(); 173 | 174 | for i in &indices { 175 | 176 | let indx = i as usize; 177 | 178 | let class = class_index_list.get(indx)?; 179 | 180 | let rect = boxes.get(indx)?; 181 | 182 | let bbox = BoxDetection{ 183 | xmin: rect.x, 184 | ymin: rect.y, 185 | xmax: rect.x + rect.width, 186 | ymax: rect.y + rect.height, 187 | conf: scores.get(indx)?, 188 | class: class 189 | }; 190 | 191 | final_boxes.push(bbox); 192 | } 193 | 194 | Ok(Detections{detections: final_boxes}) 195 | 196 | } 197 | 198 | 199 | /// Draw predicted bounding boxes. 200 | pub fn draw_predictions(img: &mut core::Mat, detections: &Detections, model_config: &ModelConfig) -> opencv::Result> { 201 | 202 | let boxes = &detections.detections; 203 | let bg_color = core::Scalar::all(255.0); 204 | let text_color = core::Scalar::all(0.0); 205 | 206 | for i in 0..boxes.len() { 207 | 208 | let bbox = &boxes[i]; 209 | let rect = opencv::core::Rect::new(bbox.xmin, bbox.ymin, bbox.xmax - bbox.xmin, bbox.ymax - bbox.ymin); 210 | 211 | let label = model_config.class_names.get(bbox.class as usize).unwrap(); 212 | 213 | opencv::imgproc::rectangle(img, rect, bg_color, 1, opencv::imgproc::LINE_8, 0)?; 214 | 215 | // draw text box above bbox 216 | let text_size = opencv::imgproc::get_text_size(&label, opencv::imgproc::FONT_HERSHEY_SIMPLEX, 0.6, 1, &mut 0).unwrap(); 217 | opencv::imgproc::rectangle(img, core::Rect{x: rect.x, y: std::cmp::max(0, rect.y - text_size.height - 2), width: rect.width, height: text_size.height + 2}, bg_color, -1, opencv::imgproc::LINE_8, 0)?; 218 | opencv::imgproc::put_text(img, &label, core::Point{x: rect.x, y: std::cmp::max(0,rect.y - 2)}, opencv::imgproc::FONT_HERSHEY_SIMPLEX, 0.6, text_color, 1, opencv::imgproc::LINE_AA, false).unwrap() 219 | 220 | } 221 | 222 | let mut out_vector :core::Vector = core::Vector::default(); 223 | opencv::imgcodecs::imencode(".jpg", img, &mut out_vector, &core::Vector::default()).unwrap(); 224 | 225 | Ok(out_vector.to_vec()) 226 | } 227 | 228 | 229 | pub fn load_model() -> Result> { 230 | 231 | let model_config = load_model_from_config().unwrap(); 232 | 233 | let model = dnn::read_net_from_onnx(&model_config.model_path); 234 | 235 | let mut model = match model { 236 | Ok(model) => model, 237 | Err(_) => {println!("Invalid ONNX model. Check out https://github.com/ultralytics/yolov5/issues/251"); std::process::exit(0)} 238 | }; 239 | model.set_preferable_backend(dnn::DNN_BACKEND_OPENCV)?; 240 | 241 | Ok(Model{model, model_config}) 242 | 243 | } 244 | 245 | /// Load model configuration. 246 | /// See ModelConfig. 247 | fn load_model_from_config() -> Result>{ 248 | 249 | let file = File::open("data/config.json"); 250 | 251 | let file = match file { 252 | Ok(file) => file, 253 | Err(_) => {println!("data/config.json does NOT exist."); std::process::exit(0)} 254 | }; 255 | 256 | let reader = BufReader::new(file); 257 | 258 | let model_config : std::result::Result = serde_json::from_reader(reader); 259 | 260 | let model_config = match model_config { 261 | Ok(model_config) => model_config, 262 | Err(_) => {println!("Malformed JSON. Check out the doc!"); std::process::exit(0)} 263 | }; 264 | 265 | if !std::path::Path::new(&model_config.model_path).exists() { 266 | println!("ONNX model in {model_path} does NOT exist.", model_path=model_config.model_path); 267 | std::process::exit(0) 268 | } 269 | 270 | println!("{model_path}", model_path=model_config.model_path); 271 | 272 | Ok(model_config) 273 | } 274 | 275 | /// Porting of letterbox padding strategy used to prepare the input image. 276 | /// 277 | /// See: https://github.com/ultralytics/yolov5/blob/master/utils/augmentations.py#L91 278 | fn letterbox( img: &core::Mat, new_shape: core::Size, scale_up: bool) -> opencv::Result { 279 | 280 | let width = img.cols() as f32; 281 | let height = img.rows() as f32; 282 | 283 | let new_width = new_shape.width as f32; 284 | let new_height = new_shape.height as f32; 285 | let mut r = f32::min(new_width / width, new_height / height ); 286 | 287 | if !scale_up { 288 | r =f32::min(r, 1.0); 289 | } 290 | 291 | let new_unpad_w = (width * r).round() as i32; 292 | let new_unpad_h = (height * r).round() as i32; 293 | 294 | let dw = (new_shape.width - new_unpad_w) / 2; 295 | let dh = (new_shape.height - new_unpad_h) / 2; 296 | 297 | let mut dst = core::Mat::default(); 298 | opencv::imgproc::resize(&img, &mut dst, core::Size_{width: new_unpad_w, height: new_unpad_h}, 0.0, 0.0, opencv::imgproc::INTER_LINEAR)?; 299 | 300 | let top = (dh as f32 - 0.1).round() as i32; 301 | let bottom = (dh as f32 + 0.1).round() as i32; 302 | let left = (dw as f32 - 0.1).round() as i32; 303 | let right = (dw as f32 + 0.1).round() as i32; 304 | 305 | let mut final_mat = core::Mat::default(); 306 | opencv::core::copy_make_border(&dst, &mut final_mat, top, bottom, left, right, opencv::core::BORDER_CONSTANT, opencv::core::Scalar::new(114.0, 114.0, 114.0, 114.0))?; 307 | 308 | //let params: core::Vector = core::Vector::default(); 309 | 310 | //opencv::imgcodecs::imwrite("padded.jpg", &final_mat, ¶ms)?; 311 | 312 | Ok(final_mat) 313 | } --------------------------------------------------------------------------------