Source

// Background removal on this device: the ISNet general model run by onnxruntime-web, on the graphics
// card through WebGPU when the browser has it, otherwise on the processor. The model and engine are
// downloaded once (see model.ts), checked against their SHA-256 and kept with the Cache API, so later
// uses work offline. Images themselves are never uploaded. Loaded only when someone turns it on.
import * as ort from "onnxruntime-web/webgpu";
import { MODEL_BYTES, MODEL_FILES } from "./model";

const SIDE = 1024;
const CACHE = "images-background-model";
let session: Promise<ort.InferenceSession> | null = null;

async function fetchBytes(url: string, report: (loaded: number) => void): Promise<Uint8Array> {
  const failed = "The background model could not be downloaded. Check your connection and try again.";
  const res = await fetch(url).catch(() => { throw new Error(failed); });
  if (!res.ok || !res.body) throw new Error(failed);
  const chunks: Uint8Array[] = [];
  let loaded = 0;
  const reader = res.body.getReader();
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    loaded += value.length;
    report(loaded);
  }
  return join(chunks);
}

function join(parts: Uint8Array[]): Uint8Array {
  const out = new Uint8Array(parts.reduce((n, p) => n + p.length, 0));
  let at = 0;
  for (const p of parts) { out.set(p, at); at += p.length; }
  return out;
}

async function sha256(bytes: Uint8Array): Promise<string> {
  const digest = new Uint8Array(await crypto.subtle.digest("SHA-256", bytes));
  return Array.from(digest, (b) => b.toString(16).padStart(2, "0")).join("");
}

/** Every file from the cache, or downloaded, checked and cached, reporting overall progress. */
async function getFiles(): Promise<Uint8Array[]> {
  // Private windows in some browsers refuse storage; then the files are downloaded each time.
  const cache = await caches.open(CACHE).catch(() => null);
  const loaded = MODEL_FILES.map(() => 0);
  let last = 0;
  const report = (i: number, n: number) => {
    loaded[i] = n;
    const sum = loaded.reduce((a, b) => a + b, 0);
    if (sum - last > 256 * 1024 || sum === MODEL_BYTES) { last = sum; self.postMessage({ progress: { loaded: sum, total: MODEL_BYTES } }); }
  };
  const files = await Promise.all(MODEL_FILES.map(async (f, i) => {
    const hit = await cache?.match(f.url);
    if (hit) {
      report(i, f.bytes);
      return new Uint8Array(await hit.arrayBuffer());
    }
    const bytes = await fetchBytes(f.url, (n) => report(i, n));
    if ((await sha256(bytes)) !== f.sha256) throw new Error("The background model did not download correctly. Try again.");
    await cache?.put(f.url, new Response(bytes)).catch(() => {}); // a full disk only means downloading again next time
    return bytes;
  }));
  // Drop files an earlier version of Images kept.
  for (const req of (await cache?.keys()) ?? []) if (!MODEL_FILES.some((f) => f.url === req.url)) await cache!.delete(req);
  return files;
}

function load(): Promise<ort.InferenceSession> {
  session ??= (async () => {
    const [wasm, model] = await getFiles();
    ort.env.wasm.wasmBinary = wasm;
    ort.env.wasm.numThreads = 1; // threads need cross-origin isolation, which the page does not have
    ort.env.logLevel = "error";
    const s = await ort.InferenceSession.create(model, { executionProviders: ["webgpu", "wasm"], logSeverityLevel: 3 });
    self.postMessage({ progress: { loaded: MODEL_BYTES, total: MODEL_BYTES } });
    return s;
  })().catch((err) => {
    session = null;
    throw err;
  });
  return session;
}

/** Returns how much of each pixel is the subject, as 1024 x 1024 values from 0 (background) to 255. */
async function cutout(src: Blob): Promise<Uint8ClampedArray> {
  const s = await load();
  const bmp = await createImageBitmap(src);
  const canvas = new OffscreenCanvas(SIDE, SIDE);
  const ctx = canvas.getContext("2d", { willReadFrequently: true })!;
  ctx.imageSmoothingQuality = "high";
  ctx.drawImage(bmp, 0, 0, SIDE, SIDE);
  bmp.close();
  const px = ctx.getImageData(0, 0, SIDE, SIDE).data;
  const n = SIDE * SIDE;
  const input = new Float32Array(3 * n);
  for (let i = 0; i < n; i++) {
    input[i] = (px[i * 4] - 128) / 256;
    input[n + i] = (px[i * 4 + 1] - 128) / 256;
    input[2 * n + i] = (px[i * 4 + 2] - 128) / 256;
  }
  const name = s.outputNames[0];
  const result = await s.run({ [s.inputNames[0]]: new ort.Tensor("float32", input, [1, 3, SIDE, SIDE]) }, [name]);
  const out = (await result[name].getData()) as Float32Array;
  result[name].dispose();
  let min = Infinity;
  let max = -Infinity;
  for (const v of out) { if (v < min) min = v; if (v > max) max = v; }
  // Stretched a little, so a sure subject is fully solid and a sure background fully clear, with no
  // faint haze left behind; edges in between (hair, fur) stay soft.
  const mask = new Uint8ClampedArray(n);
  const range = max - min || 1;
  for (let i = 0; i < n; i++) mask[i] = (((out[i] - min) / range - 0.04) / 0.92) * 255;
  return mask;
}

self.onmessage = async (e: MessageEvent<{ id: number; src: Blob }>) => {
  const { id, src } = e.data;
  try {
    const mask = await cutout(src);
    self.postMessage({ id, ok: true, value: mask }, { transfer: [mask.buffer] });
  } catch (err) {
    self.postMessage({ id, ok: false, error: (err as Error)?.message || String(err) });
  }
};