@woosh/meep-engine/src/engine/graphics/impostors/octahedral/ImpostorBaker.js

Pure JavaScript game engine. Fully featured and production ready.

import { mat4, vec3 } from "gl-matrix";
import { Mesh, OrthographicCamera, Scene, Vector4, WebGLMultipleRenderTargets } from "three";
import { assert } from "../../../../core/assert.js";
import { collectIteratorValueToArray } from "../../../../core/collection/collectIteratorValueToArray.js";
import { HashMap } from "../../../../core/collection/map/HashMap.js";
import Signal from "../../../../core/events/signal/Signal.js";
import { isPowerOfTwo } from "../../../../core/math/isPowerOfTwo.js";
import { computeMaterialEquality } from "../../../asset/loaders/material/computeMaterialEquality.js";
import { computeMaterialHash } from "../../../asset/loaders/material/computeMaterialHash.js";
import { Sampler2D } from "../../texture/sampler/Sampler2D.js";
import { three_setSceneAutoUpdate } from "../../three/three_setSceneAutoUpdate.js";
import { compute_bounding_sphere } from "./bake/compute_bounding_sphere.js";
import { prepare_bake_material } from "./bake/prepare_bake_material.js";
import { HemiOctahedralUvEncoder } from "./grid/HemiOctahedralUvEncoder.js";
import { OctahedralUvEncoder } from "./grid/OctahedralUvEncoder.js";
import { UvEncoder } from "./grid/UvEncoder.js";
import { ImpostorCaptureType } from "./ImpostorCaptureType.js";
import { ImpostorDescription } from "./ImpostorDescription.js";
import { BakeShaderStandard } from "./shader/BakeShaderStandard.js";
import { build_cutout_from_atlas_by_alpha } from "./util/build_cutout_from_atlas_by_alpha.js";

export class ImpostorBaker {

    /**
     *
     * @type {THREE.WebGLRenderer|null}
     * @private
     */
    _renderer = null;

    /**
     *
     * @param {THREE.WebGLRenderer} v
     */
    set renderer(v) {
        this._renderer = v;
    }


    /**
     *
     * @param {number[]|vec4} bounding_sphere
     * @param {{mesh:ShadedGeometry, transform:mat4}[]} objects
     * @param {number} resolution
     * @param {number} frames
     * @param {UvEncoder} encoder
     */
    bake_internal({
                      bounding_sphere,
                      objects,
                      resolution,
                      frames,
                      encoder
                  }) {

        const distance = bounding_sphere[3];

        assert.isNumber(distance, 'distance');
        assert.isNonNegativeInteger(resolution, 'resolution');
        assert.isNonNegativeInteger(frames, 'frames');

        const unit_sphere_direction = [];

        const cam = new OrthographicCamera();
        const rt = new WebGLMultipleRenderTargets(resolution, resolution, 3);

        rt.scissorTest = false;
        rt.stencilBuffer = false;

        rt.texture[0].name = 'diffuse+alpha';
        rt.texture[1].name = 'normal+depth';
        rt.texture[2].name = 'orm'; // Occlusion, Roughness, Metalness

        const renderer = this._renderer;

        // remember render state
        const _rt = renderer.getRenderTarget();
        const _vp = new Vector4();
        renderer.getViewport(_vp);
        const _autoClear = renderer.autoClear;
        const _pixelRatio = renderer.getPixelRatio();

        const gl = this._renderer.getContext();

        // Baking should be done without anti-aliasing, so we make sure to disable it
        const _enabled_antialias = gl.isEnabled(gl.SAMPLE_ALPHA_TO_COVERAGE);

        if (_enabled_antialias) {
            gl.disable(gl.SAMPLE_ALPHA_TO_COVERAGE);
        }


        // set new render state
        renderer.autoClear = false;
        renderer.setRenderTarget(rt);
        renderer.setClearColor(0xFFFFFF, 0);
        renderer.clearColor();
        renderer.setPixelRatio(1);

        const frame_width = resolution / frames;
        const frame_height = resolution / frames;

        const object_count = objects.length;

        // construct scene
        const scene = new Scene();

        three_setSceneAutoUpdate(scene, false);
        scene.matrixAutoUpdate = false;
        scene.matrixWorldNeedsUpdate = false;

        const id = new ImpostorDescription();
        id.frame_count = frames;
        id.resolution = resolution;

        /**
         * Fired when all the rendering is done to indicate that cleanup should happen
         * @type {Signal}
         */
        const cleanup_signal = new Signal();

        /**
         *
         * @type {HashMap<Material, BakeShaderStandard>}
         */
        const bake_material_map = new HashMap({
            keyHashFunction: computeMaterialHash,
            keyEqualityFunction: computeMaterialEquality
        });

        const max_anisotropic_filtering_level = renderer.capabilities.getMaxAnisotropy();

        for (let k = 0; k < object_count; k++) {
            const object = objects[k];

            const source_mesh = object.mesh;

            // ensure tangents are generated
            //buffer_geometry_ensure_tangents(source_mesh.geometry);

            const source_material = source_mesh.material;

            // obtain material for baking
            const bake_material = bake_material_map.getOrCompute(source_material, (source_material) => {

                const bake_material = new BakeShaderStandard();

                // prepare bake material to match the source material
                prepare_bake_material({
                    bake_material: bake_material,
                    source_material: source_material,
                    cleanup_signal: cleanup_signal,
                    anisotropy: max_anisotropic_filtering_level
                });

                bake_material.uniforms.uAtlasResolution.value.set(resolution, resolution);

                cleanup_signal.addOne(bake_material.dispose, bake_material);

                id.source_material_count++;

                return bake_material;
            });


            const geometry = source_mesh.geometry;
            const mesh = new Mesh(geometry, bake_material);

            mesh.matrixAutoUpdate = false;
            mesh.matrixWorldNeedsUpdate = false;
            mesh.frustumCulled = false;

            mat4.copy(mesh.matrixWorld.elements, object.transform);

            scene.add(mesh);

            // update stats
            id.source_geometry_polygon_count += geometry.getIndex().count / 3;
            id.source_geometry_vertex_count += geometry.getAttribute('position').count;
            id.source_instance_count++;
        }


        const bake_material_array = collectIteratorValueToArray([], bake_material_map.values());

        // traverse octahedron
        const max_frame_index = frames - 1;

        for (let i = 0; i < frames; i++) {
            const frame_u = max_frame_index > 0 ? (i / max_frame_index) : 0;

            for (let j = 0; j < frames; j++) {

                const frame_v = max_frame_index > 0 ? (j / max_frame_index) : 0;

                // compute vector direction where to place camera
                encoder.uv_to_unit(unit_sphere_direction, [frame_u, frame_v]);

                // offset by the radius of the sphere
                const camera_px = distance * unit_sphere_direction[0] + bounding_sphere[0];
                const camera_py = distance * unit_sphere_direction[1] + bounding_sphere[1];
                const camera_pz = distance * unit_sphere_direction[2] + bounding_sphere[2];

                // console.log(`UV:${octahedron_u.toFixed(2)},${octahedron_v.toFixed(2)}\t V3:${unit_sphere_direction.map(n => n.toFixed(2)).join(', ')}`);

                // construct projection matrix
                cam.left = -distance;
                cam.right = distance;
                cam.top = distance;
                cam.bottom = -distance;

                cam.near = 0;
                cam.far = distance * 2;

                cam.position.set(
                    camera_px,
                    camera_py,
                    camera_pz
                );

                cam.lookAt(bounding_sphere[0], bounding_sphere[1], bounding_sphere[2]);
                cam.updateProjectionMatrix();

                // update materials
                for (let k = 0; k < bake_material_array.length; k++) {
                    const mat = bake_material_array[k];

                    mat.uniforms.projection_params.value.set(
                        0, 0, 0, 1 / cam.far
                    );
                    mat.uniformsNeedUpdate = true;
                }

                //TODO consider doing super-sampling for some textures for better-looking results

                /*
                TODO it's possible to reduce bounding sphere in some cases where thin geometry would produce 0 texels
                 in the actual render, but we can't know that without doing the rendering first.
                 Consider doing a cheap pre-pass, rendering out pixels and checking if we can crop the bounding sphere further
                 */

                renderer.setViewport(i * frame_width, j * frame_height, frame_width, frame_height);
                renderer.render(scene, cam);
            }
        }

        /*
         TODO dilate non-color textures to prevent artifacts when sampling around the edges
         */

        //restore renderer state
        renderer.setRenderTarget(_rt);
        renderer.setViewport(_vp);
        renderer.autoClear = _autoClear;
        renderer.setPixelRatio(_pixelRatio);
        if (_enabled_antialias) {
            gl.enable(gl.SAMPLE_ALPHA_TO_COVERAGE);
        }

        // cleanup
        cleanup_signal.send0();

        // const sampler = convertTexture2Sampler2D(rt.texture);


        id.atlas = Sampler2D.uint8(1, resolution, resolution);
        // id.atlas = sampler;
        id.rt = rt;

        // build cutout
        id.cutout = build_cutout_from_atlas_by_alpha({
            impostor: id,
            renderer,
            vertex_limit: 5
        });


        id.sphere_radius = bounding_sphere[3];
        vec3.copy(id.offset, bounding_sphere);

        return id;
    }

    /**
     *
     * @param {{mesh:ShadedGeometry, transform:mat4}[]} objects objects that should make up the impostor. Baking will be done around the origin, so make sure to position meshes accordingly to make good use of texture space.
     * @param {number} frames how many views to capture for the atlas in each X and Y direction. Setting this to 4 will result in 16 (4*4) views captured, higher value will sacrifice detail but result in transitions being more smooth
     * @param {number} resolution resolution of the final asset, higher = more detail
     * @param {ImpostorCaptureType} type
     * @returns {ImpostorDescription}
     */
    bake({
             objects,
             frames = 12,
             resolution = 1024,
             type = ImpostorCaptureType.FullSphere
         }) {

        console.time('bake');

        assert.defined(objects, 'object');
        assert.isNonNegativeInteger(frames, 'frames');
        assert.isNonNegativeInteger(resolution, 'resolution');
        assert.greaterThanOrEqual(frames, 1, 'number of frames must be >= 1');

        if (frames > resolution) {
            throw new Error(`Frame count(=${frames}) is greater than the resolution(=${resolution}) of the bake texture, packing description is unachievable`);
        }

        const frame_resolution = resolution / frames;

        if (frames >= resolution / 4) {
            console.warn(`Bake configuration will produce less frames, this will likely result in a useless impostor that contains too little texel density to properly represent the underlying scene.  [frames=${frames}, resolution=${resolution}, resulting density = ${Math.pow(frame_resolution, 2)} texels per frame]`);
        }

        if (resolution % frames !== 0 && frame_resolution < 16) {
            console.warn(`To get results at low frame resolutions number of frames should be a divisor of resolution, instead got [frames=${frames}, resolution=${resolution}, resolution/frames=${frame_resolution}]`);
        }

        if (!isPowerOfTwo(resolution)) {
            throw new Error(` 'resolution' must be a power of two, instead was '${resolution}'`);
        }

        if (frames < 2) {
            console.warn(`Frame count is set to ${frames}, which is too low to be useful. Consider values of 2 and above. A good default is 12`);
        }

        if (this._renderer === null) {
            throw new Error('No renderer attached. Renderer is required for baking.');
        }

        const bounding_sphere = compute_bounding_sphere(objects);

        // we need to compute bounding sphere around origin of the object
        let encoder;
        if (type === ImpostorCaptureType.FullSphere) {
            encoder = new OctahedralUvEncoder();
        } else if (type === ImpostorCaptureType.Hemisphere) {
            encoder = new HemiOctahedralUvEncoder();
        } else {
            throw new Error(`Unsupported capture mode '${type}'`);
        }

        const r = this.bake_internal({
            resolution,
            frames,
            objects,
            bounding_sphere,
            encoder
        });

        r.capture_type = type;

        console.timeEnd('bake');

        return r;
    }
}