src/lib/ai/segmenter.js

import { pipeline } from "@huggingface/transformers";

/**
 * Image segmentation engine — works for both panoptic segmentation and face parsing.
 * The model ID passed to the constructor determines the behavior.
 *
 * Events:
 *   'loading'  — model download started
 *   'progress' — download progress (detail: { file, progress, loaded, total })
 *   'ready'    — model loaded and ready
 *   'error'    — something went wrong (detail: Error)
 */
export class Segmenter extends EventTarget {
  #pipe = null;
  #status = "idle";
  #modelId;
  #opts;

  /**
   * @param {string} modelId  e.g. "Xenova/detr-resnet-50-panoptic" or "jonathandinu/face-parsing"
   * @param {{ dtype?: string }} opts  Extra pipeline options (e.g. { dtype: "q8" })
   */
  constructor(modelId, opts = {}) {
    super();
    this.#modelId = modelId;
    this.#opts = opts;
  }

  get status() {
    return this.#status;
  }

  async load() {
    if (this.#status === "ready" || this.#status === "loading") return;

    this.#status = "loading";
    this.dispatchEvent(new CustomEvent("loading", { detail: { status: "loading" } }));

    try {
      const device = navigator.gpu ? "webgpu" : "wasm";
      // fp16 is ~2-3x faster on WebGPU; q8 is best for WASM CPU
      const dtype = this.#opts.dtype ?? (device === "webgpu" ? "fp16" : "q8");
      this.#pipe = await pipeline("image-segmentation", this.#modelId, {
        device,
        ...this.#opts,
        dtype,
        progress_callback: (progress) => {
          this.dispatchEvent(new CustomEvent("progress", { detail: progress }));
        },
      });
      this.#status = "ready";
      this.dispatchEvent(new CustomEvent("ready"));
    } catch (err) {
      this.#status = "error";
      this.dispatchEvent(new CustomEvent("error", { detail: err }));
      throw err;
    }
  }

  /**
   * Run segmentation on an image.
   * @param {import("@huggingface/transformers").RawImage|string} image
   * @returns {Promise<Array<{ label: string, score: number, mask: import("@huggingface/transformers").RawImage }>>}
   */
  async segment(image) {
    if (!this.#pipe) throw new Error("Segmenter model not loaded — call load() first");
    return await this.#pipe(image);
  }

  dispose() {
    this.#pipe = null;
    this.#status = "idle";
  }
}