471 lines
19 KiB
TypeScript
471 lines
19 KiB
TypeScript
/**
|
|
* Virtual Background V2 — inference worker.
|
|
*
|
|
* Handles TF.js body-segmentation (MEDIUM/HIGH tiers via WebGL/WebGPU) and TFLite WASM
|
|
* (LOW tier via selfie_segmentation_landscape) in a single dedicated Web Worker. Running
|
|
* all inference in a Worker keeps the main thread free for UI events. For GPU tiers, the
|
|
* Worker's OffscreenCanvas-backed WebGL context is not subject to Chrome's tab-visibility
|
|
* GPU scheduling throttle.
|
|
*/
|
|
/*
|
|
* Message protocol
|
|
*
|
|
* Main -> Worker:
|
|
* { type: 'init', backend: string, segWidth: number, segHeight: number,
|
|
* tfjsModelUrl: string, tfliteModelPath: string, tfliteWasmBase: string }
|
|
* { type: 'infer', bitmap: ImageBitmap } (bitmap is transferred -- zero-copy)
|
|
*
|
|
* Worker -> Main:
|
|
* { type: 'init_done', backend: string, segHeight: number, segWidth: number }
|
|
* backend/segHeight/segWidth reflect the actual tier used (may differ from requested
|
|
* backend when the Worker fell back to TFLite because the GPU backend was unavailable)
|
|
* { type: 'init_error', error: string }
|
|
* { type: 'mask', data: Uint8ClampedArray, width: number, height: number }
|
|
* (data.buffer is transferred -- zero-copy)
|
|
* { type: 'infer_error', error: string }
|
|
*/
|
|
|
|
import '@tensorflow/tfjs-backend-webgl';
|
|
import '@tensorflow/tfjs-backend-webgpu';
|
|
import * as tf from '@tensorflow/tfjs-core';
|
|
import type { BodySegmenter, MediaPipeSelfieSegmentationTfjsModelConfig }
|
|
from '@tensorflow-models/body-segmentation';
|
|
import * as bs from '@tensorflow-models/body-segmentation';
|
|
|
|
import { BackendType } from '../DeviceTierDetector';
|
|
// @ts-ignore
|
|
import createTFLiteSIMDModule from '../vendor/tflite/tflite-simd';
|
|
|
|
/** Maximum segmentation dimension in pixels. Guards against OOM via crafted init messages. */
|
|
const MAX_SEG_DIMENSION = 2048;
|
|
|
|
/**
|
|
* Extracts the effective origin of this Worker script.
|
|
*
|
|
* Workers loaded via a blob URL have self.location.href of the form
|
|
* "blob:https://hostname/uuid". Standard URL parsing returns origin="null" for blob URLs,
|
|
* so we strip the "blob:" prefix and parse the embedded URL to recover the real origin.
|
|
*
|
|
* @returns {string} The origin (e.g. "https://meet.example.com"), or empty string if unparseable.
|
|
*/
|
|
function getWorkerOrigin(): string {
|
|
try {
|
|
const { href } = self.location;
|
|
|
|
return new URL(href.startsWith('blob:') ? href.slice(5) : href).origin;
|
|
} catch {
|
|
return '';
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Returns true when {@code url} is a valid http/https URL at the same origin as this Worker.
|
|
* Rejects data URIs, javascript: URLs, file: URLs, and cross-origin paths.
|
|
*
|
|
* @param {string} url - URL to validate.
|
|
* @returns {boolean}
|
|
*/
|
|
function isSameOriginUrl(url: string): boolean {
|
|
try {
|
|
const parsed = new URL(url);
|
|
|
|
return (parsed.protocol === 'https:' || parsed.protocol === 'http:')
|
|
&& parsed.origin === getWorkerOrigin();
|
|
} catch {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/** Active TF.js segmenter (MEDIUM/HIGH tiers). Null when using TFLite backend. */
|
|
let segmenter: BodySegmenter | null = null;
|
|
|
|
/** Active TFLite module (LOW tier, TFLITE backend). Null when using TF.js. */
|
|
let tfliteModule: any = null;
|
|
|
|
/** TFLite input HEAPF32 offset (byte offset / 4). Pre-computed after model load. */
|
|
let tfliteInputOffset = 0;
|
|
|
|
/** TFLite output HEAPF32 offset (byte offset / 4). Pre-computed after model load. */
|
|
let tfliteOutputOffset = 0;
|
|
|
|
/** TFLite segmentation pixel count (segWidth x segHeight). */
|
|
let tflitePixelCount = 0;
|
|
|
|
/** TFLite segmentation canvas width. */
|
|
let tfliteSegWidth = 0;
|
|
|
|
/** TFLite segmentation canvas height. */
|
|
let tfliteSegHeight = 0;
|
|
|
|
/** OffscreenCanvas for downscaling frames before TFLite inference. */
|
|
let tfliteReadCanvas: any = null;
|
|
|
|
/** 2D context for tfliteReadCanvas. */
|
|
let tfliteReadCtx: any = null;
|
|
|
|
/** Active backend type. TFLITE selects TFLite WASM; WEBGL/WEBGPU selects TF.js. */
|
|
let currentBackend: BackendType | '' = '';
|
|
|
|
/**
|
|
* Type-safe postMessage wrapper for DedicatedWorkerGlobalScope.
|
|
* TypeScript types the global 'self' as Window (lib.dom) which has a different postMessage
|
|
* signature to DedicatedWorkerGlobalScope. Casting through 'any' bypasses the mismatch.
|
|
*
|
|
* @param {any} msg - Structured-cloneable message payload.
|
|
* @param {Transferable[]} [transfer] - Optional list of transferable objects (zero-copy).
|
|
* @returns {void}
|
|
*/
|
|
function workerPost(msg: any, transfer?: Transferable[]): void {
|
|
if (transfer?.length) {
|
|
(self as any).postMessage(msg, transfer);
|
|
} else {
|
|
(self as any).postMessage(msg);
|
|
}
|
|
}
|
|
|
|
// onmessage at module scope — runs in DedicatedWorkerGlobalScope after importScripts loads this bundle.
|
|
// TypeScript types 'onmessage' as Window.onmessage, but the runtime target is the worker global.
|
|
(self as any).onmessage = async (e: MessageEvent) => {
|
|
const { type } = e.data;
|
|
|
|
if (type === 'init') {
|
|
await handleInit(e.data);
|
|
} else if (type === 'infer') {
|
|
await handleInfer(e.data.bitmap as ImageBitmap);
|
|
}
|
|
};
|
|
|
|
/**
|
|
* Initialises the worker backend based on the tier signalled by the main thread.
|
|
*
|
|
* Routes to the TFLite handler for LOW tier (backend 'tflite') or the TF.js handler for
|
|
* MEDIUM/HIGH tiers (backend 'webgl' or 'webgpu'). The TFLite model path and WASM base are
|
|
* always passed so the TF.js handler can fall back to TFLite if GPU init fails in the Worker.
|
|
*
|
|
* @param {Object} data - Init message payload.
|
|
* @param {string} data.backend - 'tflite', 'webgl', or 'webgpu'.
|
|
* @param {number} data.segHeight - Segmentation canvas height.
|
|
* @param {number} data.segWidth - Segmentation canvas width.
|
|
* @param {string} data.tfjsModelUrl - Absolute URL of the self-hosted TF.js graph model (model.json).
|
|
* @param {string} data.tfliteModelPath - Absolute URL of the TFLite segmentation model.
|
|
* @param {string} data.tfliteWasmBase - Base URL for tflite*.wasm binaries (trailing slash).
|
|
* @returns {Promise<void>}
|
|
*/
|
|
async function handleInit(data: {
|
|
backend: BackendType;
|
|
segHeight: number;
|
|
segWidth: number;
|
|
tfjsModelUrl: string;
|
|
tfliteModelPath: string;
|
|
tfliteWasmBase: string;
|
|
}): Promise<void> {
|
|
// Whitelist backend — rejects unknown values before they reach tf.setBackend() or TFLite.
|
|
if (data.backend !== BackendType.TFLITE
|
|
&& data.backend !== BackendType.WEBGL
|
|
&& data.backend !== BackendType.WEBGPU) {
|
|
workerPost({ error: `Unknown backend: "${data.backend}"`, type: 'init_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
// Bounds-check dimensions — prevents OOM from crafted init messages.
|
|
if (!Number.isInteger(data.segWidth) || data.segWidth < 1 || data.segWidth > MAX_SEG_DIMENSION
|
|
|| !Number.isInteger(data.segHeight) || data.segHeight < 1
|
|
|| data.segHeight > MAX_SEG_DIMENSION) {
|
|
workerPost({ error: 'Invalid segmentation dimensions in init message', type: 'init_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
currentBackend = data.backend;
|
|
|
|
if (data.backend === BackendType.TFLITE) {
|
|
await handleInitTflite(data);
|
|
} else {
|
|
await handleInitTfjs(data);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Initialises TF.js with the tier-appropriate backend and creates the body-segmentation model.
|
|
* In a Worker, TF.js automatically uses OffscreenCanvas for its WebGL context (document is
|
|
* unavailable). Chrome does not apply tab-visibility GPU throttling to Worker GPU contexts.
|
|
*
|
|
* The segmentation model is loaded from the deployment's own libs/ directory. Passing an
|
|
* explicit modelUrl is required: without one, the body-segmentation library fetches the model
|
|
* from tfhub.dev, which fails in offline/firewalled deployments — and when the network hangs
|
|
* instead of erroring, the fetch outlives the main thread's init timeout so the TFLite
|
|
* fallback below never gets a chance to run.
|
|
*
|
|
* If the requested GPU backend is unavailable (no GPU, OffscreenCanvas WebGL unsupported, etc.),
|
|
* falls back to TFLite WASM using the TFLite model path supplied in the init message. This covers
|
|
* servers and VMs where the main-thread tier detector passes GPU checks (software rasteriser) but
|
|
* the Worker's OffscreenCanvas context creation fails.
|
|
*
|
|
* @param {Object} data - Init message payload.
|
|
* @param {string} data.backend - 'webgl' (MEDIUM tier) or 'webgpu' (HIGH tier).
|
|
* @param {number} data.segWidth - Unused here; inference at bitmap resolution.
|
|
* @param {number} data.segHeight - Unused here; inference at bitmap resolution.
|
|
* @param {string} data.tfjsModelUrl - Absolute URL of the self-hosted TF.js graph model (model.json).
|
|
* @param {string} data.tfliteModelPath - TFLite model URL used if GPU fallback is triggered.
|
|
* @param {string} data.tfliteWasmBase - TFLite WASM base URL used if GPU fallback is triggered.
|
|
* @returns {Promise<void>}
|
|
*/
|
|
async function handleInitTfjs(data: {
|
|
backend: string;
|
|
segHeight: number;
|
|
segWidth: number;
|
|
tfjsModelUrl: string;
|
|
tfliteModelPath: string;
|
|
tfliteWasmBase: string;
|
|
}): Promise<void> {
|
|
const { backend, tfjsModelUrl, tfliteModelPath, tfliteWasmBase } = data;
|
|
|
|
if (!isSameOriginUrl(tfjsModelUrl)) {
|
|
workerPost({ error: 'TF.js model URL must be same-origin', type: 'init_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
try {
|
|
const backendSet = await tf.setBackend(backend);
|
|
|
|
await tf.ready();
|
|
|
|
if (!backendSet || tf.getBackend() !== backend) {
|
|
// GPU backend unavailable in this Worker context — fall back to TFLite.
|
|
currentBackend = BackendType.TFLITE;
|
|
await handleInitTflite({ segHeight: 144, segWidth: 256, tfliteModelPath, tfliteWasmBase });
|
|
|
|
return;
|
|
}
|
|
|
|
segmenter = await bs.createSegmenter(
|
|
bs.SupportedModels.MediaPipeSelfieSegmentation,
|
|
{
|
|
modelType: 'landscape',
|
|
modelUrl: tfjsModelUrl,
|
|
runtime: 'tfjs'
|
|
} as MediaPipeSelfieSegmentationTfjsModelConfig
|
|
);
|
|
|
|
workerPost({ backend: currentBackend, segHeight: data.segHeight, segWidth: data.segWidth, type: 'init_done' });
|
|
} catch (err) {
|
|
// GPU init threw (e.g. OffscreenCanvas WebGL context creation failed) — fall back to TFLite.
|
|
currentBackend = BackendType.TFLITE;
|
|
try {
|
|
await handleInitTflite({ segHeight: 144, segWidth: 256, tfliteModelPath, tfliteWasmBase });
|
|
} catch (tfliteErr) {
|
|
workerPost({ error: String(tfliteErr), type: 'init_error' });
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Initialises the TFLite WASM backend for LOW tier using the selfie_segmentation_landscape model.
|
|
*
|
|
* Loads the SIMD WASM runtime via createTFLiteSIMDModule with a locateFile override so the
|
|
* .wasm binary is resolved from the deployment libs/ directory rather than from the (unusable)
|
|
* blob: worker origin. Fetches the .tflite model, writes it to the TFLite HEAPU8 buffer, and
|
|
* calls _loadModel(). Pre-computes HEAPF32 input/output offsets and creates an OffscreenCanvas
|
|
* for pixel readback.
|
|
*
|
|
* @param {Object} data - Init message payload for the TFLITE tier.
|
|
* @param {string} data.modelPath - Absolute URL of the TFLite segmentation model.
|
|
* @param {number} data.segHeight - Segmentation canvas height.
|
|
* @param {number} data.segWidth - Segmentation canvas width.
|
|
* @param {string} data.tfliteWasmBase - Base URL for tflite*.wasm binaries (trailing slash).
|
|
* @returns {Promise<void>}
|
|
*/
|
|
async function handleInitTflite(data: {
|
|
segHeight: number;
|
|
segWidth: number;
|
|
tfliteModelPath: string;
|
|
tfliteWasmBase: string;
|
|
}): Promise<void> {
|
|
try {
|
|
const { segHeight, segWidth, tfliteModelPath, tfliteWasmBase } = data;
|
|
|
|
if (!isSameOriginUrl(tfliteModelPath) || !isSameOriginUrl(`${tfliteWasmBase}tflite.wasm`)) {
|
|
workerPost({ error: 'TFLite model or WASM path must be same-origin', type: 'init_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
// locateFile overrides WASM path resolution — required because the worker runs from a
|
|
// blob: URL whose origin is "null", making relative path resolution fail.
|
|
tfliteModule = await createTFLiteSIMDModule({
|
|
locateFile: (path: string) => `${tfliteWasmBase}${path}`
|
|
});
|
|
|
|
// Fetch and load model into TFLite WASM heap.
|
|
const modelResponse = await fetch(tfliteModelPath, { credentials: 'same-origin' });
|
|
|
|
if (!modelResponse.ok) {
|
|
workerPost({ error: `TFLite model fetch failed: HTTP ${modelResponse.status}`, type: 'init_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
const modelBuffer = await modelResponse.arrayBuffer();
|
|
|
|
tfliteModule.HEAPU8.set(new Uint8Array(modelBuffer), tfliteModule._getModelBufferMemoryOffset());
|
|
tfliteModule._loadModel(modelBuffer.byteLength);
|
|
|
|
// Pre-compute HEAPF32 offsets to avoid divide-by-4 on every frame.
|
|
tfliteInputOffset = tfliteModule._getInputMemoryOffset() / 4;
|
|
tfliteOutputOffset = tfliteModule._getOutputMemoryOffset() / 4;
|
|
tfliteSegWidth = segWidth;
|
|
tfliteSegHeight = segHeight;
|
|
tflitePixelCount = segWidth * segHeight;
|
|
|
|
// Bounds-check the HEAPF32 input/output regions against the requested seg dimensions.
|
|
// Catches a mismatch between the configured seg size and the model's actual tensor
|
|
// shape before the inference loop starts silently clobbering WASM heap memory.
|
|
const heapLen = tfliteModule.HEAPF32.length;
|
|
const inputEnd = tfliteInputOffset + tflitePixelCount * 3; // NHWC float32, 3 channels.
|
|
const outputEnd = tfliteOutputOffset + tflitePixelCount; // Single float per pixel.
|
|
|
|
if (inputEnd > heapLen || outputEnd > heapLen) {
|
|
workerPost({
|
|
error: `TFLite heap bounds mismatch: segWidth=${segWidth} segHeight=${segHeight}`
|
|
+ ` requires inputEnd=${inputEnd}/outputEnd=${outputEnd} but HEAPF32 is ${heapLen}.`
|
|
+ ' Model input shape likely differs from the configured seg resolution.',
|
|
type: 'init_error'
|
|
});
|
|
|
|
return;
|
|
}
|
|
|
|
tfliteReadCanvas = new OffscreenCanvas(segWidth, segHeight);
|
|
tfliteReadCtx = tfliteReadCanvas.getContext('2d', { willReadFrequently: true });
|
|
|
|
workerPost({ backend: currentBackend, segHeight: tfliteSegHeight, segWidth: tfliteSegWidth, type: 'init_done' });
|
|
} catch (err) {
|
|
workerPost({ error: String(err), type: 'init_error' });
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Dispatches an inference call to the TFLite or TF.js handler based on the active backend.
|
|
*
|
|
* @param {ImageBitmap} bitmap - Pre-scaled camera frame at the tier's segmentation resolution.
|
|
* @returns {Promise<void>}
|
|
*/
|
|
async function handleInfer(bitmap: ImageBitmap): Promise<void> {
|
|
if (currentBackend === BackendType.TFLITE) {
|
|
await handleInferTflite(bitmap);
|
|
} else {
|
|
await handleInferTfjs(bitmap);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Runs one TF.js inference cycle on the transferred ImageBitmap.
|
|
*
|
|
* The bitmap is at the tier's segmentation resolution (pre-scaled on the main thread via
|
|
* createImageBitmap resize options). TF.js skips its own internal resize, reducing per-frame cost.
|
|
* The raw mask Uint8ClampedArray is transferred back (zero-copy) for EMA smoothing and WebGL
|
|
* compositing on the main thread.
|
|
*
|
|
* @param {ImageBitmap} bitmap - Pre-scaled camera frame at seg resolution.
|
|
* @returns {Promise<void>}
|
|
*/
|
|
async function handleInferTfjs(bitmap: ImageBitmap): Promise<void> {
|
|
if (!segmenter) {
|
|
bitmap.close();
|
|
workerPost({ error: 'Segmenter not initialised', type: 'infer_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
tf.engine().startScope();
|
|
try {
|
|
const segmentations = await segmenter.segmentPeople(
|
|
bitmap as unknown as HTMLImageElement, { flipHorizontal: false }
|
|
);
|
|
|
|
if (!segmentations.length) {
|
|
workerPost({ error: 'No segmentation result', type: 'infer_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
const maskData = await segmentations[0].mask.toImageData();
|
|
const data = maskData.data;
|
|
|
|
// Transfer the backing ArrayBuffer — zero-copy, data becomes detached in this worker.
|
|
workerPost(
|
|
{ data, height: maskData.height, type: 'mask', width: maskData.width },
|
|
[ data.buffer as ArrayBuffer ]
|
|
);
|
|
} catch (err) {
|
|
workerPost({ error: String(err), type: 'infer_error' });
|
|
} finally {
|
|
// Close unconditionally — called in finally so the bitmap is always released
|
|
// even when segmentPeople() throws before the try-block close is reached.
|
|
bitmap.close();
|
|
tf.engine().endScope();
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Runs one TFLite inference cycle on the transferred ImageBitmap.
|
|
*
|
|
* Draws the bitmap to an OffscreenCanvas, fills the TFLite HEAPF32 input buffer with NHWC
|
|
* float32 RGB values normalised to [0,1], calls _runInference(), and reads the single-channel
|
|
* output (the model outputs one float per pixel: person confidence in [0,1]).
|
|
* Writes the confidence value to both R and A channels of the returned Uint8ClampedArray so
|
|
* it is compatible with both the WebGL compositor (reads .r) and the Canvas 2D fallback (reads alpha).
|
|
*
|
|
* @param {ImageBitmap} bitmap - Pre-scaled camera frame at seg resolution.
|
|
* @returns {Promise<void>}
|
|
*/
|
|
async function handleInferTflite(bitmap: ImageBitmap): Promise<void> {
|
|
if (!tfliteModule || !tfliteReadCtx) {
|
|
bitmap.close();
|
|
workerPost({ error: 'TFLite not initialised', type: 'infer_error' });
|
|
|
|
return;
|
|
}
|
|
|
|
try {
|
|
// Draw pre-scaled bitmap to the read canvas and extract RGBA pixels.
|
|
tfliteReadCtx.drawImage(bitmap, 0, 0, tfliteSegWidth, tfliteSegHeight);
|
|
const imageData = tfliteReadCtx.getImageData(0, 0, tfliteSegWidth, tfliteSegHeight);
|
|
|
|
// Fill NHWC float32 input [1, H, W, 3] with RGB channels normalised to [0, 1].
|
|
const pixelCount = tflitePixelCount;
|
|
|
|
for (let i = 0; i < pixelCount; i++) {
|
|
tfliteModule.HEAPF32[tfliteInputOffset + i * 3] = imageData.data[i * 4] / 255;
|
|
tfliteModule.HEAPF32[tfliteInputOffset + i * 3 + 1] = imageData.data[i * 4 + 1] / 255;
|
|
tfliteModule.HEAPF32[tfliteInputOffset + i * 3 + 2] = imageData.data[i * 4 + 2] / 255;
|
|
}
|
|
|
|
// Run TFLite inference (synchronous WASM call).
|
|
tfliteModule._runInference();
|
|
|
|
// Read single-channel output: one float per pixel (person confidence in [0, 1]).
|
|
// Model outputs probabilities directly — no softmax needed.
|
|
const maskBytes = new Uint8ClampedArray(pixelCount * 4);
|
|
|
|
for (let i = 0; i < pixelCount; i++) {
|
|
const v = Math.round(tfliteModule.HEAPF32[tfliteOutputOffset + i] * 255);
|
|
|
|
maskBytes[i * 4] = v; // R — read by WebGL compositor shader (.r channel)
|
|
maskBytes[i * 4 + 3] = v; // A — read by Canvas 2D fallback compositing path
|
|
}
|
|
|
|
workerPost(
|
|
{ data: maskBytes, height: tfliteSegHeight, type: 'mask', width: tfliteSegWidth },
|
|
[ maskBytes.buffer ]
|
|
);
|
|
} catch (err) {
|
|
workerPost({ error: String(err), type: 'infer_error' });
|
|
} finally {
|
|
bitmap.close();
|
|
}
|
|
}
|
|
|