mirror of
https://github.com/huggingface/candle.git
synced 2025-06-19 03:54:56 +00:00
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:
@ -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() {}
|
||||
|
Reference in New Issue
Block a user