@livekit/track-processors 0.4.1 → 0.5.1

package/src/transformers/BackgroundTransformer.ts

@@ -5,6 +5,12 @@ import { VideoTransformerInitOptions } from './types';
 
 export type SegmenterOptions = Partial<vision.ImageSegmenterOptions['baseOptions']>;
 
+export interface FrameProcessingStats {
+  processingTimeMs: number;
+  segmentationTimeMs: number;
+  filterTimeMs: number;
+}
+
 export type BackgroundOptions = {
   blurRadius?: number;
   imagePath?: string;
@@ -12,11 +18,18 @@ export type BackgroundOptions = {
   segmenterOptions?: SegmenterOptions;
   /** cannot be updated through the `update` method, needs a restart */
   assetPaths?: { tasksVisionFileSet?: string; modelAssetPath?: string };
+  /** called when a new frame is processed */
+  onFrameProcessed?: (stats: FrameProcessingStats) => void;
 };
 
 export default class BackgroundProcessor extends VideoTransformer<BackgroundOptions> {
   static get isSupported() {
-    return typeof OffscreenCanvas !== 'undefined';
+    return (
+      typeof OffscreenCanvas !== 'undefined' &&
+      typeof VideoFrame !== 'undefined' &&
+      typeof createImageBitmap !== 'undefined' &&
+      !!document.createElement('canvas').getContext('webgl2')
+    );
   }
 
   imageSegmenter?: vision.ImageSegmenter;
@@ -25,8 +38,6 @@ export default class BackgroundProcessor extends VideoTransformer<BackgroundOpti
 
   backgroundImage: ImageBitmap | null = null;
 
-  blurRadius?: number;
-
   options: BackgroundOptions;
 
   constructor(opts: BackgroundOptions) {
@@ -36,6 +47,8 @@ export default class BackgroundProcessor extends VideoTransformer<BackgroundOpti
   }
 
   async init({ outputCanvas, inputElement: inputVideo }: VideoTransformerInitOptions) {
+    // Initialize WebGL with appropriate options based on our current state
+
     await super.init({ outputCanvas, inputElement: inputVideo });
 
     const fileSet = await vision.FilesetResolver.forVisionTasks(
@@ -51,6 +64,7 @@ export default class BackgroundProcessor extends VideoTransformer<BackgroundOpti
         delegate: 'GPU',
         ...this.options.segmenterOptions,
       },
+      canvas: this.canvas,
       runningMode: 'VIDEO',
       outputCategoryMask: true,
       outputConfidenceMasks: false,
@@ -62,6 +76,9 @@ export default class BackgroundProcessor extends VideoTransformer<BackgroundOpti
         console.error('Error while loading processor background image: ', err),
       );
     }
+    if (this.options.blurRadius) {
+      this.gl?.setBlurRadius(this.options.blurRadius);
+    }
   }
 
   async destroy() {
@@ -80,15 +97,16 @@ export default class BackgroundProcessor extends VideoTransformer<BackgroundOpti
       img.src = path;
     });
     const imageData = await createImageBitmap(img);
-    this.backgroundImage = imageData;
+    this.gl?.setBackgroundImage(imageData);
   }
 
   async transform(frame: VideoFrame, controller: TransformStreamDefaultController<VideoFrame>) {
     try {
-      if (!(frame instanceof VideoFrame)) {
+      if (!(frame instanceof VideoFrame) || frame.codedWidth === 0 || frame.codedHeight === 0) {
         console.debug('empty frame detected, ignoring');
         return;
       }
+
       if (this.isDisabled) {
         controller.enqueue(frame);
         return;
@@ -96,119 +114,53 @@ export default class BackgroundProcessor extends VideoTransformer<BackgroundOpti
       if (!this.canvas) {
         throw TypeError('Canvas needs to be initialized first');
       }
+      this.canvas.width = frame.displayWidth;
+      this.canvas.height = frame.displayHeight;
       let startTimeMs = performance.now();
-      this.imageSegmenter?.segmentForVideo(
-        this.inputVideo!,
-        startTimeMs,
-        (result) => (this.segmentationResults = result),
-      );
 
-      if (this.blurRadius) {
-        await this.blurBackground(frame);
-      } else {
-        await this.drawVirtualBackground(frame);
-      }
-      const newFrame = new VideoFrame(this.canvas, {
-        timestamp: frame.timestamp || Date.now(),
+      this.imageSegmenter?.segmentForVideo(frame, startTimeMs, (result) => {
+        const segmentationTimeMs = performance.now() - startTimeMs;
+        this.segmentationResults = result;
+        this.drawFrame(frame);
+        if (this.canvas && this.canvas.width > 0 && this.canvas.height > 0) {
+          const newFrame = new VideoFrame(this.canvas, {
+            timestamp: frame.timestamp || Date.now(),
+          });
+          const filterTimeMs = performance.now() - startTimeMs - segmentationTimeMs;
+          const stats: FrameProcessingStats = {
+            processingTimeMs: performance.now() - startTimeMs,
+            segmentationTimeMs,
+            filterTimeMs,
+          };
+          this.options.onFrameProcessed?.(stats);
+
+          controller.enqueue(newFrame);
+        } else {
+          controller.enqueue(frame);
+        }
+        frame.close();
       });
-      controller.enqueue(newFrame);
-    } finally {
+    } catch (e) {
+      console.error('Error while processing frame: ', e);
       frame?.close();
     }
   }
 
   async update(opts: BackgroundOptions) {
-    this.options = opts;
+    this.options = { ...this.options, ...opts };
     if (opts.blurRadius) {
-      this.blurRadius = opts.blurRadius;
+      this.gl?.setBlurRadius(opts.blurRadius);
     } else if (opts.imagePath) {
       await this.loadBackground(opts.imagePath);
     }
   }
 
-  async drawVirtualBackground(frame: VideoFrame) {
-    if (!this.canvas || !this.ctx || !this.segmentationResults || !this.inputVideo) return;
-    // this.ctx.save();
-    // this.ctx.clearRect(0, 0, this.canvas.width, this.canvas.height);
-    if (this.segmentationResults?.categoryMask && this.segmentationResults.categoryMask.width > 0) {
-      this.ctx.globalCompositeOperation = 'copy';
-
-      this.ctx.putImageData(
-        maskToImageData(
-          this.segmentationResults.categoryMask,
-          this.segmentationResults.categoryMask.width,
-          this.segmentationResults.categoryMask.height,
-        ),
-        0,
-        0,
-      );
-      this.ctx.filter = 'none';
-      this.ctx.globalCompositeOperation = 'source-in';
-      if (this.backgroundImage) {
-        this.ctx.drawImage(
-          this.backgroundImage,
-          0,
-          0,
-          this.backgroundImage.width,
-          this.backgroundImage.height,
-          0,
-          0,
-          this.canvas.width,
-          this.canvas.height,
-        );
-      } else {
-        this.ctx.fillStyle = '#00FF00';
-        this.ctx.fillRect(0, 0, this.canvas.width, this.canvas.height);
-      }
-
-      this.ctx.globalCompositeOperation = 'destination-over';
-    }
-    this.ctx.drawImage(frame, 0, 0, this.canvas.width, this.canvas.height);
-  }
+  async drawFrame(frame: VideoFrame) {
+    if (!this.canvas || !this.gl || !this.segmentationResults || !this.inputVideo) return;
 
-  async blurBackground(frame: VideoFrame) {
-    if (
-      !this.ctx ||
-      !this.canvas ||
-      !this.segmentationResults?.categoryMask?.canvas ||
-      !this.inputVideo
-    ) {
-      return;
+    const mask = this.segmentationResults.categoryMask;
+    if (mask) {
+      this.gl.render(frame, mask);
     }
-
-    this.ctx.save();
-    this.ctx.globalCompositeOperation = 'copy';
-
-    if (this.segmentationResults?.categoryMask && this.segmentationResults.categoryMask.width > 0) {
-      this.ctx.putImageData(
-        maskToImageData(
-          this.segmentationResults.categoryMask,
-          this.segmentationResults.categoryMask.width,
-          this.segmentationResults.categoryMask.height,
-        ),
-        0,
-        0,
-      );
-      this.ctx.filter = 'none';
-      this.ctx.globalCompositeOperation = 'source-out';
-      this.ctx.drawImage(frame, 0, 0, this.canvas.width, this.canvas.height);
-      this.ctx.globalCompositeOperation = 'destination-over';
-      this.ctx.filter = `blur(${this.blurRadius}px)`;
-      this.ctx.drawImage(frame, 0, 0, this.canvas.width, this.canvas.height);
-      this.ctx.restore();
-    }
-  }
-}
-
-function maskToImageData(mask: vision.MPMask, videoWidth: number, videoHeight: number): ImageData {
-  const dataArray: Uint8ClampedArray = new Uint8ClampedArray(videoWidth * videoHeight * 4);
-  const result = mask.getAsUint8Array();
-  for (let i = 0; i < result.length; i += 1) {
-    const offset = i * 4;
-    dataArray[offset] = result[i];
-    dataArray[offset + 1] = result[i];
-    dataArray[offset + 2] = result[i];
-    dataArray[offset + 3] = result[i];
   }
-  return new ImageData(dataArray, videoWidth, videoHeight);
 }

package/src/transformers/VideoTransformer.ts

@@ -1,3 +1,4 @@
+import { setupWebGL } from '../webgl/index';
 import { VideoTrackTransformer, VideoTransformerInitOptions } from './types';
 
 export default abstract class VideoTransformer<Options extends Record<string, unknown>>
@@ -7,10 +8,12 @@ export default abstract class VideoTransformer<Options extends Record<string, un
 
   canvas?: OffscreenCanvas;
 
-  ctx?: OffscreenCanvasRenderingContext2D;
+  // ctx?: OffscreenCanvasRenderingContext2D;
 
   inputVideo?: HTMLVideoElement;
 
+  gl?: ReturnType<typeof setupWebGL>;
+
   protected isDisabled?: Boolean = false;
 
   async init({
@@ -26,7 +29,10 @@ export default abstract class VideoTransformer<Options extends Record<string, un
     });
     this.canvas = outputCanvas || null;
     if (outputCanvas) {
-      this.ctx = this.canvas?.getContext('2d') || undefined;
+      // this.ctx = this.canvas?.getContext('2d') || undefined;
+      this.gl = setupWebGL(
+        this.canvas || new OffscreenCanvas(inputVideo.videoWidth, inputVideo.videoHeight),
+      );
     }
     this.inputVideo = inputVideo;
     this.isDisabled = false;
@@ -34,7 +40,10 @@ export default abstract class VideoTransformer<Options extends Record<string, un
 
   async restart({ outputCanvas, inputElement: inputVideo }: VideoTransformerInitOptions) {
     this.canvas = outputCanvas || null;
-    this.ctx = this.canvas.getContext('2d') || undefined;
+    this.gl?.cleanup();
+    this.gl = setupWebGL(
+      this.canvas || new OffscreenCanvas(inputVideo.videoWidth, inputVideo.videoHeight),
+    );
 
     this.inputVideo = inputVideo;
     this.isDisabled = false;
@@ -43,7 +52,8 @@ export default abstract class VideoTransformer<Options extends Record<string, un
   async destroy() {
     this.isDisabled = true;
     this.canvas = undefined;
-    this.ctx = undefined;
+    this.gl?.cleanup();
+    this.gl = undefined;
   }
 
   abstract transform(

package/src/utils.ts

@@ -1,3 +1,4 @@
+/* eslint-disable @typescript-eslint/naming-convention */
 export const supportsProcessor = typeof MediaStreamTrackGenerator !== 'undefined';
 export const supportsOffscreenCanvas = typeof OffscreenCanvas !== 'undefined';