use crate::console_log; use crate::worker::{ModelData, RunData, Worker, WorkerInput, WorkerOutput}; use wasm_bindgen::prelude::*; use wasm_bindgen_futures::JsFuture; use yew::{html, Component, Context, Html}; use yew_agent::{Bridge, Bridged}; async fn fetch_url(url: &str) -> Result, JsValue> { use web_sys::{Request, RequestCache, RequestInit, RequestMode, Response}; let window = web_sys::window().ok_or("window")?; let mut opts = RequestInit::new(); let opts = opts .method("GET") .mode(RequestMode::Cors) .cache(RequestCache::NoCache); let request = Request::new_with_str_and_init(url, opts)?; let resp_value = JsFuture::from(window.fetch_with_request(&request)).await?; // `resp_value` is a `Response` object. assert!(resp_value.is_instance_of::()); let resp: Response = resp_value.dyn_into()?; let data = JsFuture::from(resp.blob()?).await?; let blob = web_sys::Blob::from(data); let array_buffer = JsFuture::from(blob.array_buffer()).await?; let data = js_sys::Uint8Array::new(&array_buffer).to_vec(); Ok(data) } pub enum Msg { Refresh, Run, UpdateStatus(String), SetModel(ModelData), WorkerInMsg(WorkerInput), WorkerOutMsg(Result), } pub struct CurrentDecode { start_time: Option, } pub struct App { status: String, loaded: bool, generated: String, current_decode: Option, worker: Box>, } async fn model_data_load() -> Result { let weights = fetch_url("yolov8s.safetensors").await?; let model_size = "s".to_string(); console_log!("loaded weights {}", weights.len()); Ok(ModelData { weights, model_size, }) } fn performance_now() -> Option { let window = web_sys::window()?; let performance = window.performance()?; Some(performance.now() / 1000.) } fn draw_bboxes(bboxes: Vec>) -> Result<(), JsValue> { let document = web_sys::window().unwrap().document().unwrap(); let canvas = match document.get_element_by_id("canvas") { Some(canvas) => canvas, None => return Err("no canvas".into()), }; let canvas: web_sys::HtmlCanvasElement = canvas.dyn_into::()?; let context = canvas .get_context("2d")? .ok_or("no 2d")? .dyn_into::()?; let image_html_element = document.get_element_by_id("bike-img"); let image_html_element = match image_html_element { Some(data) => data, None => return Err("no bike-img".into()), }; let image_html_element = image_html_element.dyn_into::()?; canvas.set_width(image_html_element.natural_width()); canvas.set_height(image_html_element.natural_height()); context.draw_image_with_html_image_element(&image_html_element, 0., 0.)?; context.set_stroke_style(&JsValue::from("#0dff9a")); for (class_index, bboxes_for_class) in bboxes.iter().enumerate() { for b in bboxes_for_class.iter() { let name = crate::coco_classes::NAMES[class_index]; context.stroke_rect( b.xmin as f64, b.ymin as f64, (b.xmax - b.xmin) as f64, (b.ymax - b.ymin) as f64, ); if let Ok(metrics) = context.measure_text(name) { let width = metrics.width(); context.set_fill_style(&"#3c8566".into()); context.fill_rect(b.xmin as f64 - 2., b.ymin as f64 - 12., width + 4., 14.); context.set_fill_style(&"#e3fff3".into()); context.fill_text(name, b.xmin as f64, b.ymin as f64 - 2.)? } } } Ok(()) } impl Component for App { type Message = Msg; type Properties = (); fn create(ctx: &Context) -> Self { let status = "loading weights".to_string(); let cb = { let link = ctx.link().clone(); move |e| link.send_message(Self::Message::WorkerOutMsg(e)) }; let worker = Worker::bridge(std::rc::Rc::new(cb)); Self { status, generated: String::new(), current_decode: None, worker, loaded: false, } } fn rendered(&mut self, ctx: &Context, first_render: bool) { if first_render { ctx.link().send_future(async { match model_data_load().await { Err(err) => { let status = format!("{err:?}"); Msg::UpdateStatus(status) } Ok(model_data) => Msg::SetModel(model_data), } }); } } fn update(&mut self, ctx: &Context, msg: Self::Message) -> bool { match msg { Msg::SetModel(md) => { self.status = "weights loaded successfully!".to_string(); self.loaded = true; console_log!("loaded weights"); self.worker.send(WorkerInput::ModelData(md)); true } Msg::Run => { if self.current_decode.is_some() { self.status = "already processing some image at the moment".to_string() } else { let start_time = performance_now(); self.current_decode = Some(CurrentDecode { start_time }); self.status = "processing...".to_string(); self.generated.clear(); ctx.link().send_future(async { match fetch_url("bike.jpeg").await { Err(err) => { let status = format!("{err:?}"); Msg::UpdateStatus(status) } Ok(image_data) => Msg::WorkerInMsg(WorkerInput::RunData(RunData { image_data, conf_threshold: 0.5, iou_threshold: 0.5, })), } }); } true } Msg::WorkerOutMsg(output) => { match output { Ok(WorkerOutput::WeightsLoaded) => self.status = "weights loaded!".to_string(), Ok(WorkerOutput::ProcessingDone(Err(err))) => { self.status = format!("error in worker process: {err}"); self.current_decode = None } Ok(WorkerOutput::ProcessingDone(Ok(bboxes))) => { let mut content = Vec::new(); for (class_index, bboxes_for_class) in bboxes.iter().enumerate() { for b in bboxes_for_class.iter() { content.push(format!( "bbox {}: xs {:.0}-{:.0} ys {:.0}-{:.0}", crate::coco_classes::NAMES[class_index], b.xmin, b.xmax, b.ymin, b.ymax )) } } self.generated = content.join("\n"); let dt = self.current_decode.as_ref().and_then(|current_decode| { current_decode.start_time.and_then(|start_time| { performance_now().map(|stop_time| stop_time - start_time) }) }); self.status = match dt { None => "processing succeeded!".to_string(), Some(dt) => format!("processing succeeded in {:.2}s", dt,), }; self.current_decode = None; if let Err(err) = draw_bboxes(bboxes) { self.status = format!("{err:?}") } } Err(err) => { self.status = format!("error in worker {err:?}"); } } true } Msg::WorkerInMsg(inp) => { self.worker.send(inp); true } Msg::UpdateStatus(status) => { self.status = status; true } Msg::Refresh => true, } } fn view(&self, ctx: &Context) -> Html { html! {

{"Running an object detection model in the browser using rust/wasm with "} {"candle!"}

{"Once the weights have loaded, click on the run button to process an image."}

{"Source: "}{"wikimedia"}

{ if self.loaded{ html!() }else{ html! { } } }

{&self.status}

{ if self.current_decode.is_some() { html! { } } else { html! {} } }

{ self.generated.chars().map(|c| if c == '\r' || c == '\n' { html! {
} } else { html! { {c} } }).collect::() }

} } }