TmMolStar/src/extensions/mp4-export/encoder.ts

/**
 * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author David Sehnal <david.sehnal@gmail.com>
 */

import * as HME from 'h264-mp4-encoder';
import { Viewport } from '../../mol-canvas3d/camera/util';
import { ImagePass } from '../../mol-canvas3d/passes/image';
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.');
    }

    const encoder = await HME.createH264MP4Encoder();

    const { width, height } = params;
    let vw = params.viewport.width, vh = params.viewport.height;

    // dimensions must be a multiple of 2
    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;
    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, normalizedViewport);
            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.');
    }
}