[WIP] Improve Yolo WASM UI example (#591)

* return detections with classes names

* ignore .DS_Store

* example how to load wasm module

* add param to set model size

* add param for model size

* accept iou and confidence threshold on run

* conf and iou thresholds

* clamp only

* remove images from branch

* a couple of renamings, add readme with instructions

* final design

* minor font + border update

candle-wasm-examples/yolo/src/app.rs

@@ -1,5 +1,5 @@
 use crate::console_log;
-use crate::worker::{ModelData, Worker, WorkerInput, WorkerOutput};
+use crate::worker::{ModelData, RunData, Worker, WorkerInput, WorkerOutput};
 use wasm_bindgen::prelude::*;
 use wasm_bindgen_futures::JsFuture;
 use yew::{html, Component, Context, Html};
@@ -50,9 +50,13 @@ pub struct App {
 }
 
 async fn model_data_load() -> Result<ModelData, JsValue> {
-    let weights = fetch_url("yolo.safetensors").await?;
+    let weights = fetch_url("yolov8s.safetensors").await?;
+    let model_size = "s".to_string();
     console_log!("loaded weights {}", weights.len());
-    Ok(ModelData { weights })
+    Ok(ModelData {
+        weights,
+        model_size,
+    })
 }
 
 fn performance_now() -> Option<f64> {
@@ -162,7 +166,11 @@ impl Component for App {
                                 let status = format!("{err:?}");
                                 Msg::UpdateStatus(status)
                             }
-                            Ok(image_data) => Msg::WorkerInMsg(WorkerInput::Run(image_data)),
+                            Ok(image_data) => Msg::WorkerInMsg(WorkerInput::RunData(RunData {
+                                image_data,
+                                conf_threshold: 0.5,
+                                iou_threshold: 0.5,
+                            })),
                         }
                     });
                 }

candle-wasm-examples/yolo/src/bin/m.rs

@@ -1,3 +1,5 @@
+use candle_wasm_example_yolo::coco_classes;
+use candle_wasm_example_yolo::model::Bbox;
 use candle_wasm_example_yolo::worker::Model as M;
 use wasm_bindgen::prelude::*;
 
@@ -9,15 +11,36 @@ pub struct Model {
 #[wasm_bindgen]
 impl Model {
     #[wasm_bindgen(constructor)]
-    pub fn new(data: Vec<u8>) -> Result<Model, JsError> {
-        let inner = M::load_(&data)?;
+    pub fn new(data: Vec<u8>, model_size: &str) -> Result<Model, JsError> {
+        let inner = M::load_(&data, model_size)?;
         Ok(Self { inner })
     }
 
     #[wasm_bindgen]
-    pub fn run(&self, image: Vec<u8>) -> Result<String, JsError> {
-        let boxes = self.inner.run(image)?;
-        let json = serde_json::to_string(&boxes)?;
+    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 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,
+                    },
+                ));
+            }
+        }
+        let json = serde_json::to_string(&detections)?;
         Ok(json)
     }
 }

candle-wasm-examples/yolo/src/lib.rs

@@ -1,6 +1,6 @@
 mod app;
-mod coco_classes;
-mod model;
+pub mod coco_classes;
+pub mod model;
 pub mod worker;
 pub use app::App;
 pub use worker::Worker;

candle-wasm-examples/yolo/src/model.rs

@@ -623,16 +623,25 @@ fn iou(b1: &Bbox, b2: &Bbox) -> f32 {
     i_area / (b1_area + b2_area - i_area)
 }
 
-pub fn report(pred: &Tensor, img: DynamicImage, w: usize, h: usize) -> Result<Vec<Vec<Bbox>>> {
+pub fn report(
+    pred: &Tensor,
+    img: DynamicImage,
+    w: usize,
+    h: usize,
+    conf_threshold: f32,
+    iou_threshold: f32,
+) -> Result<Vec<Vec<Bbox>>> {
     let (pred_size, npreds) = pred.dims2()?;
     let nclasses = pred_size - 4;
+    let conf_threshold = conf_threshold.clamp(0.0, 1.0);
+    let iou_threshold = iou_threshold.clamp(0.0, 1.0);
     // The bounding boxes grouped by (maximum) class index.
     let mut bboxes: Vec<Vec<Bbox>> = (0..nclasses).map(|_| vec![]).collect();
     // Extract the bounding boxes for which confidence is above the threshold.
     for index in 0..npreds {
         let pred = Vec::<f32>::try_from(pred.i((.., index))?)?;
         let confidence = *pred[4..].iter().max_by(|x, y| x.total_cmp(y)).unwrap();
-        if confidence > CONFIDENCE_THRESHOLD {
+        if confidence > conf_threshold {
             let mut class_index = 0;
             for i in 0..nclasses {
                 if pred[4 + i] > pred[4 + class_index] {
@@ -659,7 +668,7 @@ pub fn report(pred: &Tensor, img: DynamicImage, w: usize, h: usize) -> Result<Ve
             let mut drop = false;
             for prev_index in 0..current_index {
                 let iou = iou(&bboxes_for_class[prev_index], &bboxes_for_class[index]);
-                if iou > NMS_THRESHOLD {
+                if iou > iou_threshold {
                     drop = true;
                     break;
                 }

candle-wasm-examples/yolo/src/worker.rs

@@ -25,6 +25,14 @@ macro_rules! console_log {
 #[derive(Serialize, Deserialize)]
 pub struct ModelData {
     pub weights: Vec<u8>,
+    pub model_size: String,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct RunData {
+    pub image_data: Vec<u8>,
+    pub conf_threshold: f32,
+    pub iou_threshold: f32,
 }
 
 pub struct Model {
@@ -32,7 +40,12 @@ pub struct Model {
 }
 
 impl Model {
-    pub fn run(&self, image_data: Vec<u8>) -> Result<Vec<Vec<Bbox>>> {
+    pub fn run(
+        &self,
+        image_data: Vec<u8>,
+        conf_threshold: f32,
+        iou_threshold: f32,
+    ) -> Result<Vec<Vec<Bbox>>> {
         console_log!("image data: {}", image_data.len());
         let image_data = std::io::Cursor::new(image_data);
         let original_image = image::io::Reader::new(image_data)
@@ -68,20 +81,37 @@ impl Model {
         let image_t = (image_t.unsqueeze(0)?.to_dtype(DType::F32)? * (1. / 255.))?;
         let predictions = self.model.forward(&image_t)?.squeeze(0)?;
         console_log!("generated predictions {predictions:?}");
-        let bboxes = report(&predictions, original_image, width, height)?;
+        let bboxes = report(
+            &predictions,
+            original_image,
+            width,
+            height,
+            conf_threshold,
+            iou_threshold,
+        )?;
         Ok(bboxes)
     }
 
-    pub fn load_(weights: &[u8]) -> Result<Self> {
+    pub fn load_(weights: &[u8], model_size: &str) -> Result<Self> {
+        let multiples = match model_size {
+            "n" => Multiples::n(),
+            "s" => Multiples::s(),
+            "m" => Multiples::m(),
+            "l" => Multiples::l(),
+            "x" => Multiples::x(),
+            _ => Err(candle::Error::Msg(
+                "invalid model size: must be n, s, m, l or x".to_string(),
+            ))?,
+        };
         let dev = &Device::Cpu;
         let weights = safetensors::tensor::SafeTensors::deserialize(weights)?;
         let vb = VarBuilder::from_safetensors(vec![weights], DType::F32, dev);
-        let model = YoloV8::load(vb, Multiples::s(), 80)?;
+        let model = YoloV8::load(vb, multiples, 80)?;
         Ok(Self { model })
     }
 
     pub fn load(md: ModelData) -> Result<Self> {
-        Self::load_(&md.weights)
+        Self::load_(&md.weights, &md.model_size.to_string())
     }
 }
 
@@ -93,7 +123,7 @@ pub struct Worker {
 #[derive(Serialize, Deserialize)]
 pub enum WorkerInput {
     ModelData(ModelData),
-    Run(Vec<u8>),
+    RunData(RunData),
 }
 
 #[derive(Serialize, Deserialize)]
@@ -125,10 +155,12 @@ impl yew_agent::Worker for Worker {
                 }
                 Err(err) => Err(format!("model creation error {err:?}")),
             },
-            WorkerInput::Run(image_data) => match &mut self.model {
+            WorkerInput::RunData(rd) => match &mut self.model {
                 None => Err("model has not been set yet".to_string()),
                 Some(model) => {
-                    let result = model.run(image_data).map_err(|e| e.to_string());
+                    let result = model
+                        .run(rd.image_data, rd.conf_threshold, rd.iou_threshold)
+                        .map_err(|e| e.to_string());
                     Ok(WorkerOutput::ProcessingDone(result))
                 }
             },