yoloWorker.js

//load the candle yolo wasm module
import init, { Model, ModelPose } from "./build/m.js";

async function fetchArrayBuffer(url) {
  const cacheName = "yolo-candle-cache";
  const cache = await caches.open(cacheName);
  const cachedResponse = await cache.match(url);
  if (cachedResponse) {
    const data = await cachedResponse.arrayBuffer();
    return new Uint8Array(data);
  }
  const res = await fetch(url, { cache: "force-cache" });
  cache.put(url, res.clone());
  return new Uint8Array(await res.arrayBuffer());
}

class Yolo {
  static instance = {};
  // Retrieve the YOLO model. When called for the first time,
  // this will load the model and save it for future use.
  static async getInstance(modelID, modelURL, modelSize) {
    // load individual modelID only once
    if (!this.instance[modelID]) {
      await init();

      self.postMessage({ status: `loading model ${modelID}:${modelSize}` });
      const weightsArrayU8 = await fetchArrayBuffer(modelURL);
      if (/pose/.test(modelID)) {
        // if pose model, use ModelPose
        this.instance[modelID] = new ModelPose(weightsArrayU8, modelSize);
      } else {
        this.instance[modelID] = new Model(weightsArrayU8, modelSize);
      }
    } else {
      self.postMessage({ status: "model already loaded" });
    }
    return this.instance[modelID];
  }
}

self.addEventListener("message", async (event) => {
  const { imageURL, modelID, modelURL, modelSize, confidence, iou_threshold } =
    event.data;
  try {
    self.postMessage({ status: "detecting" });

    const yolo = await Yolo.getInstance(modelID, modelURL, modelSize);

    self.postMessage({ status: "loading image" });
    const imgRes = await fetch(imageURL);
    const imgData = await imgRes.arrayBuffer();
    const imageArrayU8 = new Uint8Array(imgData);

    self.postMessage({ status: `running inference ${modelID}:${modelSize}` });
    const bboxes = yolo.run(imageArrayU8, confidence, iou_threshold);

    // Send the output back to the main thread as JSON
    self.postMessage({
      status: "complete",
      output: JSON.parse(bboxes),
    });
  } catch (e) {
    self.postMessage({ error: e });
  }
});