mp4 animation export wip

src/extensions/mp4-export/controls.ts

@@ -0,0 +1,144 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { debounceTime } from 'rxjs/operators';
+import { PluginStateAnimation } from '../../mol-plugin-state/animation/model';
+import { PluginComponent } from '../../mol-plugin-state/component';
+import { PluginContext } from '../../mol-plugin/context';
+import { Task } from '../../mol-task';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { encodeMp4Animation } from './encoder';
+
+export interface Mp4AnimationInfo {
+    width: number,
+    height: number
+}
+
+export const Mp4AnimationParams = {
+    quantization: PD.Numeric(18, { min: 10, max: 51 }, { description: 'Lower is better, but slower.' })
+};
+
+export class Mp4Controls extends PluginComponent {
+    private currentNames = new Set<string>();
+    private animations: PluginStateAnimation[] = [];
+
+    readonly behaviors = {
+        animations: this.ev.behavior<PD.Params>({ }),
+        current: this.ev.behavior<{ anim: PluginStateAnimation, params: PD.Params, values: any } | undefined>(void 0),
+        canApply: this.ev.behavior<PluginStateAnimation.CanApply>({ canApply: false }),
+        info: this.ev.behavior<Mp4AnimationInfo>({ width: 0, height: 0 }),
+        params: this.ev.behavior<PD.Values<typeof Mp4AnimationParams>>(PD.getDefaultValues(Mp4AnimationParams))
+    }
+
+    setCurrent(name?: string) {
+        const anim = this.animations.find(a => a.name === name);
+        if (!anim) {
+            this.behaviors.current.next(anim);
+            return;
+        }
+
+        const params = anim.params(this.plugin) as PD.Params;
+        const values = PD.getDefaultValues(params);
+
+        this.behaviors.current.next({ anim, params, values });
+        this.behaviors.canApply.next(anim.canApply?.(this.plugin) ?? { canApply: true });
+    }
+
+    setCurrentParams(values: any) {
+        this.behaviors.current.next({ ...this.behaviors.current.value!, values });
+    }
+
+    get current() {
+        return this.behaviors.current.value;
+    }
+
+    render() {
+        const task = Task.create('Export Animation', async ctx => {
+            try {
+                const resolution = this.plugin.helpers.viewportScreenshot?.getSizeAndViewport()!;
+                const anim = this.current!;
+                const movie = await encodeMp4Animation(this.plugin, ctx, {
+                    animation: {
+                        definition: anim.anim,
+                        params: anim.values,
+                    },
+                    ...resolution,
+                    quantizationParameter: this.behaviors.params.value.quantization,
+                    pass: this.plugin.helpers.viewportScreenshot?.imagePass!,
+                });
+
+                const filename = anim.anim.display.name.toLowerCase().replace(/\s/g, '-').replace(/[^a-z0-9_\-]/g, '');
+                return { movie, filename: `${this.plugin.helpers.viewportScreenshot?.getFilename('')}_${filename}.mp4` };
+            } catch (e) {
+                this.plugin.log.error('' + e);
+                throw e;
+            }
+        });
+
+        return this.plugin.runTask(task);
+    }
+
+    private get manager() {
+        return this.plugin.managers.animation;
+    }
+
+    private syncInfo() {
+        const helper = this.plugin.helpers.viewportScreenshot;
+        const size = helper?.getSizeAndViewport();
+        if (!size) return;
+
+        this.behaviors.info.next({ width: size.viewport.width, height: size.viewport.height });
+    }
+
+    private sync() {
+        const animations = this.manager.animations.filter(a => a.isExportable);
+
+        const hasAll = animations.every(a => this.currentNames.has(a.name));
+        if (hasAll && this.currentNames.size === animations.length) {
+            return;
+        }
+
+        const params = {
+            current: PD.Select(animations[0]?.name,
+                animations.map(a => [a.name, a.display.name] as [string, string]),
+                { label: 'Animation' })
+        };
+
+        const current = this.behaviors.current.value;
+        const hasCurrent = !!animations.find(a => a.name === current?.anim.name);
+
+        this.animations = animations;
+        if (!hasCurrent) {
+            this.setCurrent(animations[0]?.name);
+        }
+        this.behaviors.animations.next(params);
+    }
+
+    private init() {
+        this.subscribe(this.plugin.managers.animation.events.updated.pipe(debounceTime(16)), () => {
+            this.sync();
+        });
+
+        this.subscribe(this.plugin.canvas3d?.resized!, () => this.syncInfo());
+        this.subscribe(this.plugin.helpers.viewportScreenshot?.events.previewed!, () => this.syncInfo());
+
+        this.subscribe(this.plugin.behaviors.state.isBusy, b => {
+            const anim = this.current;
+            if (!b && anim) {
+                this.behaviors.canApply.next(anim.anim.canApply?.(this.plugin) ?? { canApply: true });
+            }
+        });
+
+        this.sync();
+        this.syncInfo();
+    }
+
+    constructor(private plugin: PluginContext) {
+        super();
+
+        this.init();
+    }
+}

src/extensions/mp4-export/encoder.ts

@@ -7,7 +7,6 @@
 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';
@@ -44,6 +43,8 @@ export async function encodeMp4Animation<A extends PluginStateAnimation>(plugin:
     if (vw % 2 !== 0) vw -= 1;
     if (vh % 2 !== 0) vh -= 1;
 
+    const normalizedViewport: Viewport = { ...params.viewport, width: vw, height: vh };
+
     encoder.width = vw;
     encoder.height = vh;
     if (params.quantizationParameter) encoder.quantizationParameter = params.quantizationParameter;
@@ -70,7 +71,7 @@ export async function encodeMp4Animation<A extends PluginStateAnimation>(plugin:
         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);
+            const image = params.pass.getImageData(width, height, normalizedViewport);
             encoder.addFrameRgba(image.data);
 
             if (ctx.shouldUpdate) {
@@ -99,89 +100,4 @@ function validateViewport(params: Mp4EncoderParams) {
     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,
-            cameraHelper: { axes: { name: 'off', params: {} } },
-            multiSample: { mode: 'on', sampleLevel: 4 }, // { mode: 'on', sampleLevel: 2 },
-            postprocessing: this.plugin.canvas3d!.props.postprocessing
-        });
-        return pass;
-    }
-
-    async generate() {
-        // const w = 1024, h = 768;
-        // const w = 1920, h = 1080;
-        const w = 1280, h = 720;
-
-        const viewport: Viewport = {
-            x: w / 2 - 300,
-            y: h / 2 - 250,
-            width: 600,
-            height: 500
-        };
-
-        const encoder = await HME.createH264MP4Encoder();
-        encoder.width = viewport.width;
-        encoder.height = viewport.height;
-        encoder.frameRate = 30;
-        encoder.quantizationParameter = 15;
-        encoder.initialize();
-
-        console.log('creating image pass');
-        const pass = this.createImagePass();
-
-        // const canvas = document.createElement('canvas');
-        // canvas.width = w;
-        // canvas.height = h;
-        // const canvasCtx = canvas.getContext('2d')!;
-
-        const loop = this.plugin.animationLoop;
-
-        loop.stop();
-        loop.resetTime(0);
-
-        const durationMs = 2500;
-        const fps = encoder.frameRate;
-
-        const color = this.plugin.canvas3d?.props.renderer.backgroundColor;
-        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, speed: 1, direction: 'cw' });
-
-        // const imageData: Uint8ClampedArray[] = [];
-        const N = Math.ceil(durationMs / 1000 * fps);
-        const dt = durationMs / N;
-        for (let i = 0; i <= N; i++) {
-            await loop.tick(i * dt, { isSynchronous: true, manualDraw: true });
-
-            const image = pass.getImageData(w, h, viewport);
-            encoder.addFrameRgba(image.data);
-            // if (i === 0) canvasCtx.putImageData(image, 0, 0);
-
-            console.log(`frame ${i + 1}/${N + 1}`);
-        }
-
-        this.plugin.canvas3d?.setProps({ renderer: { backgroundColor: color } }, true);
-
-        console.log('finalizing');
-        encoder.finalize();
-        console.log('finalized');
-        const uint8Array = encoder.FS.readFile(encoder.outputFilename);
-        console.log('encoded');
-        encoder.delete();
-
-        await this.plugin.managers.animation.stop();
-        loop.start();
-
-        return { movie: uint8Array };
-    }
-
-    constructor(private plugin: PluginContext) {
-    }
 }

src/extensions/mp4-export/ui.tsx

@@ -5,59 +5,120 @@
  */
 
 import React from 'react';
-import { AnimateCameraSpin } from '../../mol-plugin-state/animation/built-in/camera-spin';
+import { merge } from 'rxjs';
+import { debounceTime } from 'rxjs/operators';
 import { CollapsableControls, CollapsableState } from '../../mol-plugin-ui/base';
 import { Button } from '../../mol-plugin-ui/controls/common';
-import { SubscriptionsOutlinedSvg } from '../../mol-plugin-ui/controls/icons';
-import { Task } from '../../mol-task';
+import { CameraOutlinedSvg, GetAppSvg, Icon, SubscriptionsOutlinedSvg } from '../../mol-plugin-ui/controls/icons';
+import { ParameterControls } from '../../mol-plugin-ui/controls/parameters';
 import { download } from '../../mol-util/download';
-import { encodeMp4Animation, Mp4Encoder } from './encoder';
+import { Mp4AnimationParams, Mp4Controls } from './controls';
 
 interface State {
-    data?: { movie: Uint8Array };
+    busy?: boolean,
+    data?: { movie: Uint8Array, filename: string };
 }
 
 export class Mp4EncoderUI extends CollapsableControls<{}, State> {
-    protected defaultState(): State & CollapsableState  {
+    private _controls: Mp4Controls | undefined;
+
+    get controls() {
+        return this._controls || (this._controls = new Mp4Controls(this.plugin));
+    }
+
+    protected defaultState(): State & CollapsableState {
         return {
             header: 'Export Animation',
             isCollapsed: true,
             brand: { accent: 'cyan', svg: SubscriptionsOutlinedSvg }
         };
     }
-    protected renderControls(): JSX.Element | null {
+
+    private downloadControls() {
         return <>
-            <Button onClick={() => this.gen()}>Generate</Button>
-            {this.state.data && <Button onClick={() => this.save()}>Save</Button>}
+            <div className='msp-control-offset msp-help-text'>
+                <div className='msp-help-description' style={{ textAlign: 'center' }}>
+                    Rendering successful!
+                </div>
+            </div>
+            <Button icon={GetAppSvg} onClick={this.save} style={{ marginTop: 1 }}>Save Animation</Button>
+            <Button onClick={() => this.setState({ data: void 0 })} style={{ marginTop: 6 }}>Clear</Button>
         </>;
     }
 
-    save() {
-        download(new Blob([this.state.data!.movie]), 'test.mp4');
+    protected renderControls(): JSX.Element | null {
+        if (this.state.data) {
+            return this.downloadControls();
+        }
+
+        const ctrl = this.controls;
+        const current = ctrl.behaviors.current.value;
+        const info = ctrl.behaviors.info.value;
+        const canApply = ctrl.behaviors.canApply.value;
+        return <>
+            <ParameterControls
+                params={ctrl.behaviors.animations.value}
+                values={{ current: current?.anim.name }}
+                onChangeValues={xs => ctrl.setCurrent(xs.current)}
+                isDisabled={this.state.busy}
+            />
+            {current && <ParameterControls
+                params={current.params}
+                values={current.values}
+                onChangeValues={xs => ctrl.setCurrentParams(xs)}
+                isDisabled={this.state.busy}
+            />}
+            <div className='msp-control-offset msp-help-text'>
+                <div className='msp-help-description' style={{ textAlign: 'center' }}>
+                    Resolution: {info.width}x{info.height}<br />
+                    Adjust in viewport using <Icon svg={CameraOutlinedSvg} inline />
+                </div>
+            </div>
+            <ParameterControls
+                params={Mp4AnimationParams}
+                values={ctrl.behaviors.params.value}
+                onChangeValues={xs => ctrl.behaviors.params.next(xs)}
+                isDisabled={this.state.busy}
+            />
+            <Button onClick={this.generate} style={{ marginTop: 1 }}
+                disabled={this.state.busy || !canApply.canApply}
+                commit={canApply.canApply ? 'on' : 'off'}>
+                {canApply.canApply ? 'Render' : canApply.reason ?? 'Invalid params/state'}
+            </Button>
+        </>;
     }
 
-    gen() {
-        const task = Task.create('Export Animation', async ctx => {
-            const resolution = this.plugin.helpers.viewportScreenshot?.getSizeAndViewport()!;
-            const movie = await encodeMp4Animation(this.plugin, ctx, {
-                animation: {
-                    definition: AnimateCameraSpin,
-                    params: { durationInMs: 2000, speed: 1, direction: 'cw' }
-                },
-                ...resolution,
-                quantizationParameter: 18,
-                pass: this.plugin.helpers.viewportScreenshot?.imagePass!,
-            });
+    componentDidMount() {
+        const merged = merge(
+            this.controls.behaviors.animations,
+            this.controls.behaviors.current,
+            this.controls.behaviors.canApply,
+            this.controls.behaviors.info,
+            this.controls.behaviors.params
+        );
 
-            this.setState({ data: { movie } });
+        this.subscribe(merged.pipe(debounceTime(10)), () => {
+            if (!this.state.isCollapsed) this.forceUpdate();
         });
+    }
+
+    componentWillUnmount() {
+        this._controls?.dispose();
+        this._controls = void 0;
+    }
 
-        this.plugin.runTask(task);
+    save = () => {
+        download(new Blob([this.state.data!.movie]), this.state.data!.filename);
     }
 
-    async generate() {
-        const encoder = new Mp4Encoder(this.plugin);
-        const data = await encoder.generate();
-        this.setState({ data });
+    generate = async () => {
+        try {
+            this.setState({ busy: true });
+            const data = await this.controls.render();
+            console.log(data);
+            this.setState({ busy: false, data });
+        } catch {
+            this.setState({ busy: false });
+        }
     }
 }

src/mol-plugin-state/animation/built-in/assembly-unwind.ts

@@ -15,6 +15,7 @@ import { PluginContext } from '../../../mol-plugin/context';
 export const AnimateAssemblyUnwind = PluginStateAnimation.create({
     name: 'built-in.animate-assembly-unwind',
     display: { name: 'Unwind Assembly' },
+    isExportable: true,
     params: (plugin: PluginContext) => {
         const targets: [string, string][] = [['all', 'All']];
         const structures = plugin.state.data.select(StateSelection.Generators.rootsOfType(PluginStateObject.Molecule.Structure));

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

@@ -17,6 +17,7 @@ type State = { snapshot: Camera.Snapshot };
 export const AnimateCameraSpin = PluginStateAnimation.create({
     name: 'built-in.animate-camera-spin',
     display: { name: 'Camera Spin' },
+    isExportable: true,
     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.' }),

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

@@ -15,9 +15,10 @@ export { PluginStateAnimation };
 interface PluginStateAnimation<P = any, S = any> {
     name: string,
     readonly display: { readonly name: string, readonly description?: string },
+    readonly isExportable?: boolean,
 
     params(ctx: PluginContext): PD.For<P>,
-    canApply?(ctx: PluginContext): { canApply: true } | { canApply: false, reason?: string },
+    canApply?(ctx: PluginContext): PluginStateAnimation.CanApply,
     initialState(params: P, ctx: PluginContext): S,
     getDuration?(params: P, ctx: PluginContext): PluginStateAnimation.Duration,
 
@@ -39,6 +40,7 @@ interface PluginStateAnimation<P = any, S = any> {
 }
 
 namespace PluginStateAnimation {
+    export type CanApply = { canApply: true } | { canApply: false, reason?: string }
     export type Duration = { kind: 'unknown' } | { kind: 'infinite' } | { kind: 'fixed', durationMs: number  }
 
     export interface Instance<A extends PluginStateAnimation> {

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

@@ -17,7 +17,7 @@ export { PluginAnimationManager };
 
 class PluginAnimationManager extends StatefulPluginComponent<PluginAnimationManager.State> {
     private map = new Map<string, PluginStateAnimation>();
-    private animations: PluginStateAnimation[] = [];
+    private _animations: PluginStateAnimation[] = [];
     private currentTime: number = 0;
 
     private _current: PluginAnimationManager.Current;
@@ -28,9 +28,11 @@ class PluginAnimationManager extends StatefulPluginComponent<PluginAnimationMana
         applied: this.ev(),
     };
 
-    get isEmpty() { return this.animations.length === 0; }
+    get isEmpty() { return this._animations.length === 0; }
     get current() { return this._current!; }
 
+    get animations() { return this._animations; }
+
     private triggerUpdate() {
         this.events.updated.next();
     }
@@ -42,8 +44,8 @@ class PluginAnimationManager extends StatefulPluginComponent<PluginAnimationMana
     getParams(): PD.Params {
         if (!this._params) {
             this._params = {
-                current: PD.Select(this.animations[0] && this.animations[0].name,
-                    this.animations.map(a => [a.name, a.display.name] as [string, string]),
+                current: PD.Select(this._animations[0] && this._animations[0].name,
+                    this._animations.map(a => [a.name, a.display.name] as [string, string]),
                     { label: 'Animation' })
             };
         }
@@ -79,8 +81,8 @@ class PluginAnimationManager extends StatefulPluginComponent<PluginAnimationMana
         }
         this._params = void 0;
         this.map.set(animation.name, animation);
-        this.animations.push(animation);
-        if (this.animations.length === 1) {
+        this._animations.push(animation);
+        if (this._animations.length === 1) {
             this.updateParams({ current: animation.name });
         } else {
             this.triggerUpdate();

src/mol-plugin/index.ts

@@ -88,8 +88,8 @@ export const DefaultPluginSpec: PluginSpec = {
     ],
     animations: [
         AnimateModelIndex,
-        AnimateAssemblyUnwind,
-        AnimateCameraSpin
+        AnimateCameraSpin,
+        AnimateAssemblyUnwind
     ]
 };
 

src/mol-plugin/util/viewport-screenshot.ts

@@ -68,7 +68,11 @@ class ViewportScreenshotHelper extends PluginComponent {
             resolution: this.params.resolution.defaultValue
         }),
         cropParams: this.ev.behavior<{ auto: boolean, relativePadding: number }>({ auto: true, relativePadding: 0.1 }),
-        relativeCrop: this.ev.behavior<Viewport>({ x: 0, y: 0, width: 1, height: 1 })
+        relativeCrop: this.ev.behavior<Viewport>({ x: 0, y: 0, width: 1, height: 1 }),
+    };
+
+    readonly events = {
+        previewed: this.ev<any>()
     };
 
     get values() {
@@ -121,15 +125,24 @@ class ViewportScreenshotHelper extends PluginComponent {
 
     private _imagePass: ImagePass;
     get imagePass() {
-        return this._imagePass || (this._imagePass = this.createPass(true));
+        if (this._imagePass) {
+            this._imagePass.setProps({
+                cameraHelper: { axes: this.values.axes },
+                transparentBackground: this.values.transparent,
+                // TODO: optimize because this creates a copy of a large object!
+                postprocessing: this.plugin.canvas3d!.props.postprocessing
+            });
+            return this._imagePass;
+        }
+        return this._imagePass = this.createPass(true);
     }
 
-    getFilename() {
+    getFilename(extension = '.png') {
         const models = this.plugin.state.data.select(StateSelection.Generators.rootsOfType(PluginStateObject.Molecule.Model)).map(s => s.obj!.data);
         const uniqueIds = new Set<string>();
         models.forEach(m => uniqueIds.add(m.entryId.toUpperCase()));
         const idString = SetUtils.toArray(uniqueIds).join('-');
-        return `${idString || 'molstar-image'}.png`;
+        return `${idString || 'molstar-image'}${extension}`;
     }
 
     private canvas = function () {
@@ -250,6 +263,8 @@ class ViewportScreenshotHelper extends PluginComponent {
         if (!canvasCtx) throw new Error('Could not create canvas 2d context');
         canvasCtx.putImageData(imageData, 0, 0);
         if (this.cropParams.auto) this.autocrop();
+
+        this.events.previewed.next();
         return { canvas, width: w, height: h };
     }
 
@@ -272,13 +287,6 @@ class ViewportScreenshotHelper extends PluginComponent {
         if (width <= 0 || height <= 0) return;
 
         await ctx.update('Rendering image...');
-        this.imagePass.setProps({
-            cameraHelper: { axes: this.values.axes },
-            transparentBackground: this.values.transparent,
-            // TODO: optimize because this creates a copy of a large object!
-            postprocessing: this.plugin.canvas3d!.props.postprocessing
-        });
-
         const imageData = this.imagePass.getImageData(width, height, viewport);
 
         await ctx.update('Encoding image...');