Add wasm support for yolo-v8 pose detection. (#630)

* Add wasm support for yolo-v8 pose detection.

* Better bbox handling.

* Add the pose model in the wasm example lib.
This commit is contained in:
Laurent Mazare
2023-08-27 19:49:24 +01:00
committed by GitHub
parent 72ebb12bca
commit 24dda44c27
3 changed files with 315 additions and 40 deletions

View File

@ -1,6 +1,7 @@
use candle_wasm_example_yolo::coco_classes;
use candle_wasm_example_yolo::model::Bbox;
use candle_wasm_example_yolo::worker::Model as M;
use candle_wasm_example_yolo::worker::ModelPose as P;
use wasm_bindgen::prelude::*;
#[wasm_bindgen]
@ -26,18 +27,9 @@ impl Model {
let bboxes = self.inner.run(image, conf_threshold, iou_threshold)?;
let mut detections: Vec<(String, Bbox)> = vec![];
for (class_index, bboxes_for_class) in bboxes.iter().enumerate() {
for b in bboxes_for_class.iter() {
detections.push((
coco_classes::NAMES[class_index].to_string(),
Bbox {
xmin: b.xmin,
ymin: b.ymin,
xmax: b.xmax,
ymax: b.ymax,
confidence: b.confidence,
},
));
for (class_index, bboxes_for_class) in bboxes.into_iter().enumerate() {
for b in bboxes_for_class.into_iter() {
detections.push((coco_classes::NAMES[class_index].to_string(), b));
}
}
let json = serde_json::to_string(&detections)?;
@ -45,4 +37,30 @@ impl Model {
}
}
#[wasm_bindgen]
pub struct ModelPose {
inner: P,
}
#[wasm_bindgen]
impl ModelPose {
#[wasm_bindgen(constructor)]
pub fn new(data: Vec<u8>, model_size: &str) -> Result<ModelPose, JsError> {
let inner = P::load_(&data, model_size)?;
Ok(Self { inner })
}
#[wasm_bindgen]
pub fn run(
&self,
image: Vec<u8>,
conf_threshold: f32,
iou_threshold: f32,
) -> Result<String, JsError> {
let bboxes = self.inner.run(image, conf_threshold, iou_threshold)?;
let json = serde_json::to_string(&bboxes)?;
Ok(json)
}
}
fn main() {}