refactor(web): extract JS results payload into payload module (#249)

Signed-off-by: Onuralp SEZER <onuralp@ultralytics.com>

Author: onuralpszr

crates/web/src/lib.rs

@@ -4,7 +4,7 @@
 //!
 //! Runs YOLO ONNX models in the browser on WebGPU via ONNX Runtime Web (bridged
 //! by [`ort-web`](https://ort.pyke.io/backends/web)), reusing the core crate's
-//! preprocessing/postprocessing so results match the native and Python paths.
+//! preprocessing/postprocessing so results match the native path.
 //! The published npm package wraps this behind a small `YOLO` class. The whole
 //! crate is gated to `wasm32`; elsewhere it compiles to an empty library.
 #![cfg(target_arch = "wasm32")]
@@ -26,14 +26,17 @@ use ultralytics_inference::postprocessing::postprocess;
 use ultralytics_inference::preprocessing::{
     preprocess_image_center_crop, preprocess_image_with_precision,
 };
-use ultralytics_inference::results::{Results, SemanticMask, Speed};
+use ultralytics_inference::results::Speed;
 use ultralytics_inference::visualizer::color::Color;
 use ultralytics_inference::visualizer::skeleton::{
     KPT_COLOR_INDICES, LIMB_COLOR_INDICES, SKELETON,
 };
 use ultralytics_inference::{InferenceConfig, Task};
 
+use payload::JsResults;
+
 mod onnx_meta;
+mod payload;
 
 /// Default inference image size used when a model does not record `imgsz` in its
 /// metadata. Mirrors the native crate's fallback.
@@ -206,8 +209,7 @@ impl YoloModel {
     ///
     /// Initializes the accelerated backend on first use, reads the embedded
     /// Ultralytics metadata from the model bytes, then commits an ONNX Runtime
-    /// session on the requested execution provider. Equivalent to Python's
-    /// `YOLO('model.onnx')`.
+    /// session on the requested execution provider.
     ///
     /// `device` is `"webgpu"` or `"cpu"` (the browser picks the GPU adapter
     /// automatically). WebGPU is registered with `error_on_failure`, so if the
@@ -452,236 +454,9 @@ fn to_js<T: Serialize>(value: &T, what: &str) -> Result<JsValue, JsError> {
         .map_err(|e| JsError::new(&format!("failed to serialize {what}: {e}")))
 }
 
-/// One detected box, mirroring `Boxes` in the Ultralytics API (pixel `xyxy`).
-/// `color` is the Ultralytics palette color for the class (`#rrggbb`).
-#[derive(Serialize)]
-struct JsBox {
-    x1: f32,
-    y1: f32,
-    x2: f32,
-    y2: f32,
-    conf: f32,
-    cls: usize,
-    name: String,
-    color: String,
-}
-
-/// One oriented box (`xywhr` + score), mirroring `Obb`.
-#[derive(Serialize)]
-struct JsObb {
-    x: f32,
-    y: f32,
-    w: f32,
-    h: f32,
-    angle: f32,
-    conf: f32,
-    cls: usize,
-    name: String,
-    color: String,
-}
-
-/// Keypoints for one detection (`[[x, y, conf], ...]`), with the instance color.
-#[derive(Serialize)]
-struct JsKeypoints {
-    points: Vec<[f32; 3]>,
-    color: String,
-}
-
-/// Per-stage timing in ms, mirroring `Speed` (decode+letterbox / inference+sync / decode+NMS).
-#[derive(Serialize)]
-struct JsSpeed {
-    preprocess: f64,
-    inference: f64,
-    postprocess: f64,
-}
-
-/// Classification probabilities, mirroring `Probs`.
-#[derive(Serialize)]
-struct JsProbs {
-    top1: usize,
-    top5: Vec<usize>,
-    top1conf: f32,
-    top5conf: Vec<f32>,
-    name: String,
-    color: String,
-}
-
-/// JS-facing results payload, mirroring the Ultralytics `Results` API. Per-task
-/// fields are empty/`null` when unused.
-#[derive(Serialize)]
-struct JsResults {
-    task: String,
-    width: u32,
-    height: u32,
-    boxes: Vec<JsBox>,
-    obb: Vec<JsObb>,
-    keypoints: Vec<JsKeypoints>,
-    probs: Option<JsProbs>,
-    /// Segment masks as a translucent colored `RGBA` overlay (`width*height*4`);
-    /// empty for other tasks. A `Uint8Array`, drawable straight onto a canvas.
-    #[serde(with = "serde_bytes")]
-    masks: Vec<u8>,
-    /// Semantic class id per pixel, little-endian `u16` (`width*height*2` bytes),
-    /// row-major; empty for other tasks. The TS wrapper reinterprets it as a
-    /// `Uint16Array`. The IGNORE sentinel `65535` marks background or
-    /// class-filtered pixels.
-    #[serde(with = "serde_bytes")]
-    semantic: Vec<u8>,
-    /// Per-stage timing in ms: `{ preprocess, inference, postprocess }`.
-    speed: JsSpeed,
-}
-
-impl JsResults {
-    /// Convert core [`Results`] into the serializable JS payload, labeling it with
-    /// the model's declared `task` (the caller already holds it, so there is no
-    /// need to guess it back from which result fields are populated).
-    fn from_results(r: &Results, task: Task) -> Self {
-        // Class id -> display name and palette color, shared by every detection type.
-        let name = |c: usize| r.names.get(&c).cloned().unwrap_or_default();
-        let hex = |c: usize| Color::from_index(c).to_hex();
-
-        let boxes = r.boxes.as_ref().map_or_else(Vec::new, |b| {
-            let (xyxy, conf, cls) = (b.xyxy(), b.conf(), b.cls());
-            (0..b.len())
-                .map(|i| {
-                    let c = cls[i] as usize;
-                    JsBox {
-                        x1: xyxy[[i, 0]],
-                        y1: xyxy[[i, 1]],
-                        x2: xyxy[[i, 2]],
-                        y2: xyxy[[i, 3]],
-                        conf: conf[i],
-                        cls: c,
-                        name: name(c),
-                        color: hex(c),
-                    }
-                })
-                .collect()
-        });
-
-        let obb = r.obb.as_ref().map_or_else(Vec::new, |o| {
-            let (xywhr, conf, cls) = (o.xywhr(), o.conf(), o.cls());
-            (0..conf.len())
-                .map(|i| {
-                    let c = cls[i] as usize;
-                    JsObb {
-                        x: xywhr[[i, 0]],
-                        y: xywhr[[i, 1]],
-                        w: xywhr[[i, 2]],
-                        h: xywhr[[i, 3]],
-                        angle: xywhr[[i, 4]],
-                        conf: conf[i],
-                        cls: c,
-                        name: name(c),
-                        color: hex(c),
-                    }
-                })
-                .collect()
-        });
-
-        let keypoints = r.keypoints.as_ref().map_or_else(Vec::new, |k| {
-            let (n, npt) = (k.data.shape()[0], k.data.shape()[1]);
-            (0..n)
-                .map(|i| JsKeypoints {
-                    points: (0..npt)
-                        .map(|j| [k.data[[i, j, 0]], k.data[[i, j, 1]], k.data[[i, j, 2]]])
-                        .collect(),
-                    // Skeleton uses the matching detection's color.
-                    color: boxes.get(i).map_or_else(|| hex(0), |b| b.color.clone()),
-                })
-                .collect()
-        });
-
-        let probs = r.probs.as_ref().map(|p| JsProbs {
-            top1: p.top1(),
-            top5: p.top5(),
-            top1conf: p.top1conf(),
-            top5conf: p.top5conf(),
-            name: name(p.top1()),
-            color: hex(p.top1()),
-        });
-
-        Self {
-            task: task.as_str().to_owned(),
-            width: r.orig_shape.1,
-            height: r.orig_shape.0,
-            boxes,
-            obb,
-            keypoints,
-            probs,
-            masks: build_mask_overlay(r),
-            semantic: r.semantic_mask.as_ref().map_or_else(Vec::new, |sem| {
-                let mut bytes = Vec::with_capacity(sem.data.len() * 2);
-                for &v in &sem.data {
-                    bytes.extend_from_slice(&v.to_le_bytes());
-                }
-                bytes
-            }),
-            speed: JsSpeed {
-                preprocess: r.speed.preprocess.unwrap_or(0.0),
-                inference: r.speed.inference.unwrap_or(0.0),
-                postprocess: r.speed.postprocess.unwrap_or(0.0),
-            },
-        }
-    }
-}
-
-/// Write one colored RGBA pixel at `(x, y)` into a `width`-wide buffer.
-fn put(buf: &mut [u8], w: usize, x: usize, y: usize, c: Color, a: u8) {
-    let i = (y * w + x) * 4;
-    buf[i..i + 4].copy_from_slice(&[c.0, c.1, c.2, a]);
-}
-
-/// Build a translucent colored RGBA overlay (`width*height*4`) from the segment
-/// instance masks or the semantic class map, colored by the class palette.
-/// Empty for other tasks or if the mask resolution does not match the image.
-fn build_mask_overlay(r: &Results) -> Vec<u8> {
-    let (h, w) = (r.orig_shape.0 as usize, r.orig_shape.1 as usize);
-    let mut buf = vec![0u8; w * h * 4];
-
-    // Segment: per-instance binary masks, each colored by its detection class.
-    if let (Some(masks), Some(boxes)) = (&r.masks, &r.boxes) {
-        let (n, mh, mw) = masks.data.dim();
-        if mh != h || mw != w {
-            return Vec::new();
-        }
-        let cls = boxes.cls();
-        for i in 0..n.min(cls.len()) {
-            let c = Color::from_index(cls[i] as usize);
-            for y in 0..h {
-                for x in 0..w {
-                    if masks.data[[i, y, x]] > 0.5 {
-                        put(&mut buf, w, x, y, c, 140);
-                    }
-                }
-            }
-        }
-        return buf;
-    }
-
-    // Semantic: per-pixel class map, blended 50/50 like the native renderer.
-    if let Some(sem) = &r.semantic_mask {
-        if sem.data.dim() != (h, w) {
-            return Vec::new();
-        }
-        for y in 0..h {
-            for x in 0..w {
-                // Filtered-out classes carry the IGNORE sentinel; leave transparent.
-                let cls = sem.data[[y, x]];
-                if cls != SemanticMask::IGNORE {
-                    put(&mut buf, w, x, y, Color::from_index(cls as usize), 128);
-                }
-            }
-        }
-        return buf;
-    }
-
-    Vec::new()
-}
-
 /// The Ultralytics pose drawing scheme: skeleton connectivity plus the per-limb
 /// and per-keypoint palette colors. Lets the JS annotator draw pose exactly like
-/// the native/Python renderer without duplicating any palette.
+/// the native renderer without duplicating any palette.
 #[derive(Serialize)]
 struct PoseScheme {
     /// Keypoint index pairs connected by limbs.