mp4 encoder wip

src/extensions/mp4-export/encoder.ts

@@ -1,10 +1,100 @@
 import * as HME from 'h264-mp4-encoder';
 import { Viewport } from '../../mol-canvas3d/camera/util';
+import { ImagePass } from '../../mol-canvas3d/passes/image';
 import { AnimateCameraSpin } from '../../mol-plugin-state/animation/built-in/camera-spin';
+import { PluginStateAnimation } from '../../mol-plugin-state/animation/model';
 import { PluginContext } from '../../mol-plugin/context';
+import { RuntimeContext } from '../../mol-task';
+import { Color } from '../../mol-util/color';
+
+
+export interface Mp4EncoderParams<A extends PluginStateAnimation = PluginStateAnimation> {
+    pass: ImagePass,
+    customBackground?: Color,
+    animation: PluginStateAnimation.Instance<A>,
+    width: number,
+    height: number,
+    viewport?: Viewport,
+    /** default is 30 */
+    fps?: number,
+    /** Number from 10 (best quality, slowest) to 51 (worst, fastest) */
+    quantizationParameter?: number
+}
+
+export async function encodeMp4Animation<A extends PluginStateAnimation>(plugin: PluginContext, ctx: RuntimeContext, params: Mp4EncoderParams<A>) {
+    await ctx.update({ message: 'Initializing...', isIndeterminate: true });
+
+    validateViewport(params);
+    const durationMs = PluginStateAnimation.getDuration(plugin, params.animation);
+    if (durationMs === void 0) {
+        throw new Error('The animation does not have the duration specified.');
+    }
 
-export class Mp4Encoder {
+    const encoder = await HME.createH264MP4Encoder();
+
+    const { width, height } = params;
+    const vw = params.viewport?.width ?? width, vh = params.viewport?.height ?? height;
+
+
+    encoder.width = vw;
+    encoder.height = vh;
+    if (params.quantizationParameter) encoder.quantizationParameter = params.quantizationParameter;
+    if (params.fps) encoder.frameRate = params.fps;
+    encoder.initialize();
+
+    const loop = plugin.animationLoop;
+    const originalBackground = params.customBackground ? plugin.canvas3d?.props.renderer.backgroundColor : void 0;
+    let stoppedAnimation = true, finalized = false;
+
+    try {
+        loop.stop();
+        loop.resetTime(0);
+        plugin.canvas3d?.setProps({ renderer: { backgroundColor: params.customBackground } }, true);
+
+        const fps = encoder.frameRate;
+        const N = Math.ceil(durationMs / 1000 * fps);
+        const dt = durationMs / N;
+
+        await ctx.update({ message: 'Rendering...', isIndeterminate: false, current: 0, max: N + 1 });
+
+        await plugin.managers.animation.play(params.animation.definition, params.animation.params);
+        stoppedAnimation = false;
+        for (let i = 0; i <= N; i++) {
+            await loop.tick(i * dt, { isSynchronous: true, manualDraw: true });
+
+            const image = params.pass.getImageData(width, height, params.viewport);
+            encoder.addFrameRgba(image.data);
+
+            if (ctx.shouldUpdate) {
+                await ctx.update({ current: i + 1 });
+            }
+        }
+        await ctx.update({ message: 'Applying finishing touches...', isIndeterminate: true });
+        await plugin.managers.animation.stop();
+        stoppedAnimation = true;
+        encoder.finalize();
+        finalized = true;
+        return encoder.FS.readFile(encoder.outputFilename);
+    } finally {
+        if (finalized) encoder.delete();
+        if (originalBackground) {
+            plugin.canvas3d?.setProps({ renderer: { backgroundColor: originalBackground } }, true);
+        }
+        if (!stoppedAnimation) await plugin.managers.animation.stop();
+        loop.start();
+    }
+}
+
+function validateViewport(params: Mp4EncoderParams) {
+    if (!params.viewport) return;
+
+    if (params.viewport.x + params.viewport.width > params.width || params.viewport.x + params.viewport.width >= params.width) {
+        throw new Error('Viewport exceeds the canvas dimensions.');
+    }
+}
 
+
+export class Mp4Encoder {
     createImagePass() {
         const pass = this.plugin.canvas3d!.getImagePass({
             transparentBackground: true,
@@ -15,10 +105,6 @@ export class Mp4Encoder {
         return pass;
     }
 
-    sleep() {
-        return new Promise(res => setTimeout(res, 16.6));
-    }
-
     async generate() {
         // const w = 1024, h = 768;
         // const w = 1920, h = 1080;
@@ -58,7 +144,7 @@ export class Mp4Encoder {
         this.plugin.canvas3d?.setProps({ renderer: { backgroundColor: 0x000000 as any } }, true);
 
         // await this.plugin.managers.animation.play(AnimateAssemblyUnwind, { durationInMs: 3000, playOnce: true, target: 'all' });
-        await this.plugin.managers.animation.play(AnimateCameraSpin, { durationInMs: durationMs, direction: 'cw' });
+        await this.plugin.managers.animation.play(AnimateCameraSpin, { durationInMs: durationMs, speed: 1, direction: 'cw' });
 
         // const imageData: Uint8ClampedArray[] = [];
         const N = Math.ceil(durationMs / 1000 * fps);

src/extensions/mp4-export/ui.tsx

@@ -1,8 +1,10 @@
 import React from 'react';
+import { AnimateCameraSpin } from '../../mol-plugin-state/animation/built-in/camera-spin';
 import { CollapsableControls, CollapsableState } from '../../mol-plugin-ui/base';
 import { Button } from '../../mol-plugin-ui/controls/common';
+import { Task } from '../../mol-task';
 import { download } from '../../mol-util/download';
-import { Mp4Encoder } from './encoder';
+import { encodeMp4Animation, Mp4Encoder } from './encoder';
 
 interface State {
     data?: { movie: Uint8Array };
@@ -18,7 +20,7 @@ export class Mp4EncoderTestUI extends CollapsableControls<{}, State> {
     }
     protected renderControls(): JSX.Element | null {
         return <>
-            <Button onClick={() => this.generate()}>Generate</Button>
+            <Button onClick={() => this.gen()}>Generate</Button>
             {this.state.data && <Button onClick={() => this.save()}>Save</Button>}
         </>;
     }
@@ -28,6 +30,25 @@ export class Mp4EncoderTestUI extends CollapsableControls<{}, State> {
         // download(this.state.data!.image, 'test.png');
     }
 
+    gen() {
+        const task = Task.create('Export Animation', async ctx => {
+            const movie = await encodeMp4Animation(this.plugin, ctx, {
+                animation: {
+                    definition: AnimateCameraSpin,
+                    params: { durationInMs: 2000, speed: 1, direction: 'cw' }
+                },
+                width: 1280,
+                height: 720,
+                quantizationParameter: 18,
+                pass: this.plugin.helpers.viewportScreenshot?.imagePass!,
+            });
+
+            this.setState({ data: { movie } });
+        });
+
+        this.plugin.runTask(task);
+    }
+
     async generate() {
         const encoder = new Mp4Encoder(this.plugin);
         const data = await encoder.generate();

src/mol-plugin-state/animation/built-in/camera-spin.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  */
@@ -19,9 +19,11 @@ export const AnimateCameraSpin = PluginStateAnimation.create({
     display: { name: 'Camera Spin' },
     params: () => ({
         durationInMs: PD.Numeric(4000, { min: 100, max: 20000, step: 100 }),
+        speed: PD.Numeric(1, { min: 1, max: 10, step: 1 }, { description: 'How many times to spin in the specified dutation.' }),
         direction: PD.Select<'cw' | 'ccw'>('cw', [['cw', 'Clockwise'], ['ccw', 'Counter Clockwise']], { cycle: true })
     }),
     initialState: (_, ctx) => ({ snapshot: ctx.canvas3d?.camera.getSnapshot()! }) as State,
+    getDuration: p => ({ kind: 'fixed', durationMs: p.durationInMs }),
     teardown: (_, state: State, ctx) => {
         ctx.canvas3d?.requestCameraReset({ snapshot: state.snapshot, durationMs: 0 });
     },
@@ -42,7 +44,7 @@ export const AnimateCameraSpin = PluginStateAnimation.create({
             return { kind: 'finished' };
         }
 
-        const angle = 2 * Math.PI * phase * (ctx.params.direction === 'ccw' ? -1 : 1);
+        const angle = 2 * Math.PI * phase * ctx.params.speed * (ctx.params.direction === 'ccw' ? -1 : 1);
 
         Vec3.sub(_dir, snapshot.position, snapshot.target);
         Vec3.normalize(_axis, snapshot.up);

src/mol-plugin-state/animation/model.ts

@@ -19,6 +19,7 @@ interface PluginStateAnimation<P = any, S = any> {
     params(ctx: PluginContext): PD.For<P>,
     canApply?(ctx: PluginContext): { canApply: true } | { canApply: false, reason?: string },
     initialState(params: P, ctx: PluginContext): S,
+    getDuration?(params: P, ctx: PluginContext): PluginStateAnimation.Duration,
 
     // TODO: support state in setup/teardown?
     setup?(params: P, ctx: PluginContext): void | Promise<void>,
@@ -38,6 +39,14 @@ interface PluginStateAnimation<P = any, S = any> {
 }
 
 namespace PluginStateAnimation {
+    export type Duration = { kind: 'unknown' } | { kind: 'infinite' } | { kind: 'fixed', durationMs: number  }
+
+    export interface Instance<A extends PluginStateAnimation> {
+        definition: PluginStateAnimation,
+        params: Params<A>,
+        customDurationMs?: number
+    }
+
     export interface Time {
         lastApplied: number,
         current: number
@@ -49,7 +58,15 @@ namespace PluginStateAnimation {
         plugin: PluginContext
     }
 
+    export type Params<A extends PluginStateAnimation> = A extends PluginStateAnimation<infer P> ? P : never
+
     export function create<P, S>(params: PluginStateAnimation<P, S>) {
         return params;
     }
+
+    export function getDuration<A extends PluginStateAnimation>(ctx: PluginContext, instance: Instance<A>) {
+        if (instance.customDurationMs) return instance.customDurationMs;
+        const d = instance.definition.getDuration?.(instance.params, ctx);
+        if (d?.kind === 'fixed') return d.durationMs;
+    }
 }

src/mol-plugin-state/manager/animation.ts

@@ -125,7 +125,7 @@ class PluginAnimationManager extends StatefulPluginComponent<PluginAnimationMana
 
         this._current.lastTime = 0;
         this._current.startedTime = -1;
-        this._current.state = this._current.anim.initialState(anim, this.context);
+        this._current.state = this._current.anim.initialState(this._current.paramValues, this.context);
         this.isStopped = false;
     }