242 lines
9.6 KiB
TypeScript
242 lines
9.6 KiB
TypeScript
import { getBaseUrl } from '../../../base/util/helpers';
|
|
import logger from '../../../virtual-background/logger';
|
|
import { BackendType, IDeviceCapabilities } from '../DeviceTierDetector';
|
|
|
|
/**
|
|
* Worker-based segmentation backend.
|
|
*
|
|
* Wraps the {@link VBInferenceWorker} message protocol. The worker runs TF.js body-segmentation
|
|
* (MEDIUM/HIGH tiers via WebGL/WebGPU) or TFLite WASM (LOW tier via selfie_segmentation_landscape)
|
|
* in a dedicated Web Worker thread.
|
|
*
|
|
* The worker may fall back from a GPU backend to TFLite if the GPU context is unavailable inside
|
|
* the Worker. When this happens, the capabilities are updated to reflect the actual backend and
|
|
* segmentation dimensions.
|
|
*
|
|
* EMA temporal smoothing is NOT applied here — that is the processor's responsibility. The backend
|
|
* returns raw mask ImageData from each inference call.
|
|
*/
|
|
export default class WorkerSegmentationBackend {
|
|
_capabilities: IDeviceCapabilities;
|
|
_pendingInferHandler: ((e: MessageEvent) => void) | null = null;
|
|
_pendingInferResolve: ((value: ImageData | null) => void) | null = null;
|
|
_worker: Worker | null = null;
|
|
_workerReady = false;
|
|
|
|
/**
|
|
* Creates a new worker-based segmentation backend.
|
|
*
|
|
* @param {IDeviceCapabilities} capabilities - Detected device capabilities.
|
|
*/
|
|
constructor(capabilities: IDeviceCapabilities) {
|
|
this._capabilities = { ...capabilities };
|
|
}
|
|
|
|
/**
|
|
* Returns the device capabilities for this backend.
|
|
*
|
|
* @returns {IDeviceCapabilities} The capabilities (may be updated after worker fallback).
|
|
*/
|
|
get capabilities(): IDeviceCapabilities {
|
|
return this._capabilities;
|
|
}
|
|
|
|
/**
|
|
* Creates the inference worker and sends the init message.
|
|
*
|
|
* The worker is loaded via an importScripts blob URL pattern (same approach as
|
|
* FaceLandmarksDetector) to load the pre-built worker bundle from /libs/ without a separate
|
|
* module-worker instantiation.
|
|
*
|
|
* Resolves when the worker signals init_done. If the worker falls back to a different backend,
|
|
* capabilities are updated to reflect the actual backend and segmentation dimensions.
|
|
*
|
|
* @returns {Promise<void>} Resolves when the worker is ready for inference.
|
|
* @throws {Error} When the worker signals init_error.
|
|
*/
|
|
async init(): Promise<void> {
|
|
const base = getBaseUrl().replace(/\/?$/, '/');
|
|
const workerUrl = `${base}libs/vb-inference-worker.min.js`;
|
|
|
|
// JSON.stringify escapes quotes, backslashes, and control chars in workerUrl.
|
|
const scriptSource = `importScripts(${JSON.stringify(workerUrl)});`;
|
|
|
|
// @ts-ignore — @types/node's BlobOptions / URL shadow DOM BlobPropertyBag / URL here.
|
|
const blob = new Blob([ scriptSource ], { type: 'application/javascript' });
|
|
|
|
// @ts-ignore
|
|
const blobUrl = window.URL.createObjectURL(blob);
|
|
|
|
this._worker = new Worker(blobUrl, { name: 'VB V2 Inference' });
|
|
|
|
// @ts-ignore
|
|
window.URL.revokeObjectURL(blobUrl);
|
|
|
|
const { promise, resolve, reject } = Promise.withResolvers<void>();
|
|
|
|
const handler = (e: MessageEvent) => {
|
|
if (e.data.type === 'init_done') {
|
|
this._worker?.removeEventListener('message', handler);
|
|
|
|
// Worker must report which backend it actually loaded. A missing or unknown
|
|
// value means the init_done payload is malformed and capabilities cannot be
|
|
// trusted (the rest of the pipeline keys off _capabilities.backend) — reject
|
|
// so the failure surfaces instead of silently running on the wrong backend.
|
|
const reportedBackend = e.data.backend as string;
|
|
const validBackends: string[] = Object.values(BackendType);
|
|
|
|
if (!reportedBackend || !validBackends.includes(reportedBackend)) {
|
|
reject(new Error(
|
|
'[WorkerBackend] init_done missing or invalid backend:'
|
|
+ ` ${reportedBackend ?? '<missing>'}`
|
|
));
|
|
|
|
return;
|
|
}
|
|
|
|
this._workerReady = true;
|
|
|
|
// Worker may have fallen back to TFLite if the GPU backend was unavailable.
|
|
if (reportedBackend !== this._capabilities.backend) {
|
|
logger.debug(
|
|
`[WorkerBackend] Worker fell back from ${this._capabilities.backend}`
|
|
+ ` to ${reportedBackend}`
|
|
);
|
|
this._capabilities = {
|
|
...this._capabilities,
|
|
backend: reportedBackend as BackendType,
|
|
segHeight: e.data.segHeight ?? this._capabilities.segHeight,
|
|
segWidth: e.data.segWidth ?? this._capabilities.segWidth
|
|
};
|
|
}
|
|
|
|
logger.debug(
|
|
`[WorkerBackend] Inference worker ready (tier: ${this._capabilities.tier},`
|
|
+ ` backend: ${this._capabilities.backend})`
|
|
);
|
|
resolve();
|
|
} else if (e.data.type === 'init_error') {
|
|
this._worker?.removeEventListener('message', handler);
|
|
reject(new Error(`Inference worker init error: ${e.data.error}`));
|
|
}
|
|
};
|
|
|
|
this._worker.addEventListener('message', handler);
|
|
|
|
const tfliteModelPath = `${base}libs/selfie_segmentation_landscape.tflite`;
|
|
|
|
// Without an explicit URL the body-segmentation library fetches from tfhub.dev,
|
|
// breaking offline/firewalled deployments. Deployed alongside TFLite via Makefile.
|
|
const tfjsModelUrl = `${base}libs/selfie_segmentation_landscape_tfjs/model.json`;
|
|
|
|
logger.debug(
|
|
`[WorkerBackend] Initialising inference worker — backend: ${this._capabilities.backend}`
|
|
+ ` | model: ${tfliteModelPath}`
|
|
+ ` | seg: ${this._capabilities.segWidth}x${this._capabilities.segHeight}`
|
|
);
|
|
|
|
this._worker.postMessage({
|
|
backend: this._capabilities.backend,
|
|
segHeight: this._capabilities.segHeight,
|
|
segWidth: this._capabilities.segWidth,
|
|
tfjsModelUrl,
|
|
tfliteModelPath,
|
|
tfliteWasmBase: `${base}libs/`,
|
|
type: 'init'
|
|
});
|
|
|
|
const timeout = new Promise<never>((_, rej) => {
|
|
setTimeout(() => rej(new Error('Inference worker init timed out (30s)')), 30_000);
|
|
});
|
|
|
|
return Promise.race([ promise, timeout ]);
|
|
}
|
|
|
|
/**
|
|
* Sends a pre-scaled ImageBitmap to the inference worker and awaits the mask.
|
|
*
|
|
* The bitmap is transferred (zero-copy). The worker runs segmentation and returns the raw
|
|
* mask Uint8ClampedArray (also transferred). The mask is wrapped in an ImageData and returned
|
|
* directly — no EMA smoothing is applied here.
|
|
*
|
|
* Only one inference call can be in flight at a time. The caller (processor) is responsible
|
|
* for enforcing sequential call discipline.
|
|
*
|
|
* @param {ImageBitmap} bitmap - Pre-scaled camera frame at seg resolution.
|
|
* @returns {Promise<ImageData | null>} Raw mask, or null if inference failed.
|
|
*/
|
|
infer(bitmap: ImageBitmap): Promise<ImageData | null> {
|
|
if (!this._worker || !this._workerReady) {
|
|
bitmap.close();
|
|
|
|
return Promise.resolve(null);
|
|
}
|
|
|
|
const { promise, resolve } = Promise.withResolvers<ImageData | null>();
|
|
|
|
this._pendingInferResolve = resolve;
|
|
|
|
const handler = (e: MessageEvent) => {
|
|
if (e.data.type === 'mask') {
|
|
this._worker?.removeEventListener('message', handler);
|
|
if (this._pendingInferHandler === handler) {
|
|
this._pendingInferHandler = null;
|
|
}
|
|
if (this._pendingInferResolve === resolve) {
|
|
this._pendingInferResolve = null;
|
|
}
|
|
|
|
resolve(new ImageData(
|
|
e.data.data as Uint8ClampedArray,
|
|
e.data.width as number,
|
|
e.data.height as number
|
|
));
|
|
} else if (e.data.type === 'infer_error') {
|
|
this._worker?.removeEventListener('message', handler);
|
|
if (this._pendingInferHandler === handler) {
|
|
this._pendingInferHandler = null;
|
|
}
|
|
if (this._pendingInferResolve === resolve) {
|
|
this._pendingInferResolve = null;
|
|
}
|
|
logger.warn('[WorkerBackend] Worker inference error:', e.data.error);
|
|
resolve(null);
|
|
}
|
|
};
|
|
|
|
this._pendingInferHandler = handler;
|
|
this._worker.addEventListener('message', handler);
|
|
this._worker.postMessage(
|
|
{ bitmap, type: 'infer' },
|
|
[ bitmap as unknown as Transferable ]
|
|
);
|
|
|
|
return promise;
|
|
}
|
|
|
|
/**
|
|
* Stops the inference worker and releases all resources.
|
|
*
|
|
* Aborts any pending inference call (removes the message handler and resolves with null),
|
|
* sends a stop message to the worker, and terminates it.
|
|
*
|
|
* @returns {Promise<void>}
|
|
*/
|
|
async stop(): Promise<void> {
|
|
if (this._pendingInferHandler && this._worker) {
|
|
this._worker.removeEventListener('message', this._pendingInferHandler);
|
|
}
|
|
this._pendingInferHandler = null;
|
|
this._pendingInferResolve?.(null);
|
|
this._pendingInferResolve = null;
|
|
|
|
if (this._worker) {
|
|
this._worker.terminate();
|
|
this._worker = null;
|
|
}
|
|
this._workerReady = false;
|
|
|
|
logger.debug('[WorkerBackend] Stopped');
|
|
}
|
|
}
|