[web-wasm] Fix for calculating frame buffer size

software-mansion · Feb 14, 2025 · f1d6115 · f1d6115
1 parent ee7bf0e
commit f1d6115
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 39 deletions.
diff --git a/ts/smelter-web-wasm/src/workerContext/input/MediaStreamInput.ts b/ts/smelter-web-wasm/src/workerContext/input/MediaStreamInput.ts
@@ -4,6 +4,7 @@ import { InputVideoFrameRef } from './frame';
 import type { Interval } from '../../utils';
 import { SmelterEventType } from '../../eventSender';
 import { workerPostEvent } from '../bridge';
+import type { Logger } from 'pino';
 
 export type InputState = 'started' | 'playing' | 'finished';
 
@@ -18,9 +19,12 @@ export class MediaStreamInput implements Input {
   private sentEos: boolean = false;
   private sentFirstFrame: boolean = false;
 
-  public constructor(inputId: InputId, source: ReadableStream) {
+  private logger: Logger;
+
+  public constructor(inputId: InputId, source: ReadableStream, logger: Logger) {
     this.reader = source.getReader();
     this.inputId = inputId;
+    this.logger = logger;
   }
 
   public start(): InputStartResult {
@@ -35,10 +39,13 @@ export class MediaStreamInput implements Input {
         if (this.frameRef) {
           this.frameRef.decrementRefCount();
         }
-        this.frameRef = new InputVideoFrameRef({
-          frame: readResult.value,
-          ptsMs: 0, // pts does not matter here
-        });
+        this.frameRef = new InputVideoFrameRef(
+          {
+            frame: readResult.value,
+            ptsMs: 0, // pts does not matter here
+          },
+          this.logger
+        );
       }
 
       if (readResult.done) {

diff --git a/ts/smelter-web-wasm/src/workerContext/input/QueuedInput.ts b/ts/smelter-web-wasm/src/workerContext/input/QueuedInput.ts
@@ -137,7 +137,7 @@ export class QueuedInput implements Input {
       this.firstFramePtsMs = frame.ptsMs;
     }
     frame.ptsMs = frame.ptsMs - this.firstFramePtsMs;
-    return new InputVideoFrameRef(frame);
+    return new InputVideoFrameRef(frame, this.logger);
   }
 
   /**

diff --git a/ts/smelter-web-wasm/src/workerContext/input/frame.ts b/ts/smelter-web-wasm/src/workerContext/input/frame.ts
@@ -1,6 +1,7 @@
 import type { Frame } from '@swmansion/smelter-browser-render';
 import { FrameFormat } from '@swmansion/smelter-browser-render';
 import { assert } from '../../utils';
+import type { Logger } from 'pino';
 
 export type InputVideoFrame = {
   frame: Omit<VideoFrame, 'timestamp'>;
@@ -17,9 +18,11 @@ export class InputVideoFrameRef {
   private frame: InputVideoFrame;
   private refCount: number;
   private downloadedFrame?: Frame;
+  private logger: Logger;
 
-  public constructor(frame: InputVideoFrame) {
+  public constructor(frame: InputVideoFrame, logger: Logger) {
     this.frame = frame;
+    this.logger = logger;
     this.refCount = 1;
   }
 
@@ -56,43 +59,79 @@ export class InputVideoFrameRef {
     assert(this.refCount > 0);
 
     if (!this.downloadedFrame) {
-      this.downloadedFrame = await downloadFrame(this.frame);
+      this.downloadedFrame = await this.downloadFrame(this.frame);
     }
     return this.downloadedFrame;
   }
-}
 
-export class NonCopyableFrameRef extends InputVideoFrameRef {
-  public constructor(frame: InputVideoFrame) {
-    super(frame);
-  }
+  private async downloadFrame(inputFrame: InputVideoFrame): Promise<Frame> {
+    // TODO: Add support back from safari
+    // Safari does not support conversion to RGBA
+    // Chrome does not support conversion to YUV
 
-  public incrementRefCount(): void {
-    throw new Error('Reference count of `NonCopyableFrameRef` cannot be incremented');
+    const frame = inputFrame.frame;
+
+    // visibleRect is undefined when inputFrame is detached
+    assert(frame.visibleRect);
+
+    const options = {
+      format: 'RGBA',
+      layout: [
+        {
+          offset: 0,
+          stride: frame.visibleRect.width * 4,
+        },
+      ],
+    };
+
+    const buffer = new Uint8ClampedArray(frame.allocationSize(options as VideoFrameCopyToOptions));
+    const planeLayouts = await frame.copyTo(buffer, options as VideoFrameCopyToOptions);
+
+    if (!checkPlaneLayouts(options.layout, planeLayouts)) {
+      const frameInfo = {
+        displayWidth: frame.displayWidth,
+        displayHeight: frame.displayHeight,
+        codedWidth: frame.codedWidth,
+        codedHeight: frame.codedHeight,
+        visibleRect: frame.visibleRect,
+        codedRect: frame.codedRect,
+        format: frame.format,
+        colorSpace: frame.colorSpace,
+        duration: frame.duration,
+      };
+
+      this.logger.error(
+        { planeLayouts, frameInfo },
+        "Copied frame's plane layouts do not match expected layouts"
+      );
+    }
+
+    return {
+      resolution: {
+        width: frame.visibleRect.width,
+        height: frame.visibleRect.height,
+      },
+      format: FrameFormat.RGBA_BYTES,
+      data: buffer,
+    };
   }
 }
 
-async function downloadFrame(inputFrame: InputVideoFrame): Promise<Frame> {
-  // Safari does not support conversion to RGBA
-  // Chrome does not support conversion to YUV
-
-  // TODO: Add support back from safari
-  // const isSafari = !!(window as any).safari;
-  const isSafari = false;
-  const options = {
-    format: isSafari ? 'I420' : 'RGBA',
-  };
-
-  const frame = inputFrame.frame;
-  const buffer = new Uint8ClampedArray(frame.allocationSize(options as VideoFrameCopyToOptions));
-  await frame.copyTo(buffer, options as VideoFrameCopyToOptions);
-
-  return {
-    resolution: {
-      width: frame.displayWidth,
-      height: frame.displayHeight,
-    },
-    format: isSafari ? FrameFormat.YUV_BYTES : FrameFormat.RGBA_BYTES,
-    data: buffer,
-  };
+/**
+ * Returns `true` if plane layouts are valid
+ */
+function checkPlaneLayouts(expected: PlaneLayout[], received: PlaneLayout[]): boolean {
+  if (expected.length !== received.length) {
+    return false;
+  }
+  for (let i = 0; i < expected.length; i++) {
+    if (expected[i].offset !== received[i].offset) {
+      return false;
+    }
+    if (expected[i].stride !== received[i].stride) {
+      return false;
+    }
+  }
+
+  return true;
 }
diff --git a/ts/smelter-web-wasm/src/workerContext/input/input.ts b/ts/smelter-web-wasm/src/workerContext/input/input.ts
@@ -62,7 +62,7 @@ export async function createInput(
     return new QueuedInput(inputId, source, inputLogger);
   } else if (request.type === 'stream') {
     assert(request.videoStream);
-    return new MediaStreamInput(inputId, request.videoStream);
+    return new MediaStreamInput(inputId, request.videoStream, inputLogger);
   }
   throw new Error(`Unknown input type ${(request as any).type}`);
 }