VolumePass

Source

index.js

import macro from 'vtk.js/Sources/macros';
import vtkPolyData from 'vtk.js/Sources/Common/DataModel/PolyData';
import vtkProperty from 'vtk.js/Sources/Rendering/Core/Property';
import vtkRenderPass from 'vtk.js/Sources/Rendering/SceneGraph/RenderPass';
import vtkWebGPUBufferManager from 'vtk.js/Sources/Rendering/WebGPU/BufferManager';
import vtkWebGPUSimpleMapper from 'vtk.js/Sources/Rendering/WebGPU/SimpleMapper';
import vtkWebGPURenderEncoder from 'vtk.js/Sources/Rendering/WebGPU/RenderEncoder';
import vtkWebGPUShaderCache from 'vtk.js/Sources/Rendering/WebGPU/ShaderCache';
import vtkWebGPUTexture from 'vtk.js/Sources/Rendering/WebGPU/Texture';
import vtkWebGPUUniformBuffer from 'vtk.js/Sources/Rendering/WebGPU/UniformBuffer';
import vtkWebGPUFullScreenQuad from 'vtk.js/Sources/Rendering/WebGPU/FullScreenQuad';
import vtkWebGPUVolumePassFSQ from 'vtk.js/Sources/Rendering/WebGPU/VolumePassFSQ';
import * as vtkMath from 'vtk.js/Sources/Common/Core/Math';

const { Representation } = vtkProperty;
const { BufferUsage, PrimitiveTypes } = vtkWebGPUBufferManager;

// The volume rendering pass consists of two sub passes. The first
// (depthRange) renders polygonal cubes for the volumes to compute min and
// max bounds in depth for the image. This is then fed into the second pass
// (final) which actually does the raycasting between those bounds sampling
// the volumes along the way. So the first pass tends to be very fast whicle
// the second is where most of the work is done.

// given x then y then z ordering
//
//     2-----3
//   / |   / |
//  6-----7  |
//  |  |  |  |
//  |  0-----1
//  |/    |/
//  4-----5
//
const cubeFaceTriangles = [
  [0, 4, 6],
  [0, 6, 2],
  [1, 3, 7],
  [1, 7, 5],
  [0, 5, 4],
  [0, 1, 5],
  [2, 6, 7],
  [2, 7, 3],
  [0, 3, 1],
  [0, 2, 3],
  [4, 5, 7],
  [4, 7, 6],
];

const DepthBoundsFS = `
//VTK::Renderer::Dec

//VTK::Select::Dec

//VTK::VolumePass::Dec

//VTK::TCoord::Dec

//VTK::RenderEncoder::Dec

//VTK::Mapper::Dec

//VTK::IOStructs::Dec

@fragment
fn main(
//VTK::IOStructs::Input
)
//VTK::IOStructs::Output
{
  var output : fragmentOutput;

  //VTK::Select::Impl

  //VTK::TCoord::Impl

  //VTK::VolumePass::Impl

  // use the maximum (closest) of the current value and the zbuffer
  // the blend func will then take the min to find the farthest stop value
  var stopval: f32 = max(input.fragPos.z, textureLoad(opaquePassDepthTexture, vec2<i32>(i32(input.fragPos.x), i32(input.fragPos.y)), 0));

  //VTK::RenderEncoder::Impl
  return output;
}
`;

const volumeCopyFragTemplate = `
//VTK::Renderer::Dec

//VTK::Mapper::Dec

//VTK::TCoord::Dec

//VTK::RenderEncoder::Dec

//VTK::IOStructs::Dec

@fragment
fn main(
//VTK::IOStructs::Input
)
//VTK::IOStructs::Output
{
  var output: fragmentOutput;

  var computedColor: vec4<f32> = textureSample(volumePassColorTexture,
    volumePassColorTextureSampler, mapperUBO.tscale*input.tcoordVS);

  //VTK::RenderEncoder::Impl
  return output;
}
`;

/* eslint-disable no-undef */
/* eslint-disable no-bitwise */

// ----------------------------------------------------------------------------

function vtkWebGPUVolumePass(publicAPI, model) {
  // Set our className
  model.classHierarchy.push('vtkWebGPUVolumePass');

  // create the required textures, encoders, FSQ etc
  publicAPI.initialize = (viewNode) => {
    if (!model._clearEncoder) {
      publicAPI.createClearEncoder(viewNode);
    }

    if (!model._mergeEncoder) {
      publicAPI.createMergeEncoder(viewNode);
    }

    if (!model._copyEncoder) {
      publicAPI.createCopyEncoder(viewNode);
    }

    if (!model._depthRangeEncoder) {
      publicAPI.createDepthRangeEncoder(viewNode);
    }

    if (!model.fullScreenQuad) {
      model.fullScreenQuad = vtkWebGPUVolumePassFSQ.newInstance();
      model.fullScreenQuad.setDevice(viewNode.getDevice());
      model.fullScreenQuad.setTextureViews([
        ...model._depthRangeEncoder.getColorTextureViews(),
      ]);
    }

    if (!model._volumeCopyQuad) {
      model._volumeCopyQuad = vtkWebGPUFullScreenQuad.newInstance();
      model._volumeCopyQuad.setPipelineHash('volpassfsq');
      model._volumeCopyQuad.setDevice(viewNode.getDevice());
      model._volumeCopyQuad.setFragmentShaderTemplate(volumeCopyFragTemplate);
      model._copyUBO = vtkWebGPUUniformBuffer.newInstance({
        label: 'mapperUBO',
      });
      model._copyUBO.addEntry('tscale', 'vec2<f32>');
      model._volumeCopyQuad.setUBO(model._copyUBO);
      model._volumeCopyQuad.setTextureViews([model._colorTextureView]);
    }
  };

  publicAPI.traverse = (renNode, viewNode) => {
    if (model.deleted) {
      return;
    }

    // we just render our delegates in order
    model._currentParent = viewNode;

    // create stuff we need
    publicAPI.initialize(viewNode);

    // determine if we are rendering a small size
    publicAPI.computeTiming(viewNode);

    // first render the boxes to generate a min max depth
    // map for all the volumes
    publicAPI.renderDepthBounds(renNode, viewNode);

    // always mark true
    model._firstGroup = true;

    const device = viewNode.getDevice();

    // determine how many volumes we can render at a time. We subtract
    // 4 because we use know we use textures for min, max, ofun and tfun
    const maxVolumes =
      device.getHandle().limits.maxSampledTexturesPerShaderStage - 4;

    // if we have to make multiple passes then break the volumes up into groups
    // rendered from farthest to closest
    if (model.volumes.length > maxVolumes) {
      const cameraPos = renNode.getRenderable().getActiveCamera().getPosition();
      // sort from back to front based on volume centroid
      const distances = [];
      for (let v = 0; v < model.volumes.length; v++) {
        const bounds = model.volumes[v].getRenderable().getBounds();
        const centroid = [
          0.5 * (bounds[1] + bounds[0]),
          0.5 * (bounds[3] + bounds[2]),
          0.5 * (bounds[5] + bounds[4]),
        ];
        distances[v] = vtkMath.distance2BetweenPoints(centroid, cameraPos);
      }

      // sort by distance
      const volumeOrder = [...Array(model.volumes.length).keys()];
      volumeOrder.sort((a, b) => distances[b] - distances[a]);

      // render in chunks back to front
      let volumesToRender = [];
      // start with smallest chunk so that the last (closest) chunk
      // has a full maxVolumes;
      let chunkSize = volumeOrder.length % maxVolumes;
      for (let v = 0; v < volumeOrder.length; v++) {
        volumesToRender.push(model.volumes[volumeOrder[v]]);
        if (volumesToRender.length >= chunkSize) {
          publicAPI.rayCastPass(viewNode, renNode, volumesToRender);
          volumesToRender = [];
          chunkSize = maxVolumes;
          model._firstGroup = false;
        }
      }
    } else {
      // if not rendering in chunks then just draw all of them at once
      publicAPI.rayCastPass(viewNode, renNode, model.volumes);
    }

    // copy back to the original color buffer

    // final composite
    model._volumeCopyQuad.setWebGPURenderer(renNode);
    if (model._useSmallViewport) {
      const width = model._colorTextureView.getTexture().getWidth();
      const height = model._colorTextureView.getTexture().getHeight();
      model._copyUBO.setArray('tscale', [
        model._smallViewportWidth / width,
        model._smallViewportHeight / height,
      ]);
    } else {
      model._copyUBO.setArray('tscale', [1.0, 1.0]);
    }
    model._copyUBO.sendIfNeeded(device);

    model._copyEncoder.setColorTextureView(0, model.colorTextureView);
    model._copyEncoder.attachTextureViews();

    model._copyEncoder.begin(viewNode.getCommandEncoder());
    renNode.scissorAndViewport(model._copyEncoder);
    model._volumeCopyQuad.prepareAndDraw(model._copyEncoder);
    model._copyEncoder.end();
  };

  // unsubscribe from our listeners
  publicAPI.delete = macro.chain(() => {
    if (model._animationRateSubscription) {
      model._animationRateSubscription.unsubscribe();
      model._animationRateSubscription = null;
    }
  }, publicAPI.delete);

  publicAPI.computeTiming = (viewNode) => {
    const rwi = viewNode.getRenderable().getInteractor();

    if (model._lastScale == null) {
      const firstMapper = model.volumes[0].getRenderable().getMapper();
      model._lastScale = firstMapper.getInitialInteractionScale() || 1.0;
    }
    model._useSmallViewport = false;
    if (rwi.isAnimating() && model._lastScale > 1.5) {
      model._useSmallViewport = true;
    }

    model._colorTexture.resize(
      viewNode.getCanvas().width,
      viewNode.getCanvas().height
    );

    if (!model._animationRateSubscription) {
      // when the animation frame rate changes recompute the scale factor
      model._animationRateSubscription = rwi.onAnimationFrameRateUpdate(() => {
        const firstMapper = model.volumes[0].getRenderable().getMapper();
        if (firstMapper.getAutoAdjustSampleDistances()) {
          const frate = rwi.getRecentAnimationFrameRate();
          const targetScale =
            (model._lastScale * rwi.getDesiredUpdateRate()) / frate;

          model._lastScale = targetScale;
          // clamp scale to some reasonable values.
          // Below 1.5 we will just be using full resolution as that is close enough
          // Above 400 seems like a lot so we limit to that 1/20th per axis
          if (model._lastScale > 400) {
            model._lastScale = 400;
          }
        } else {
          model._lastScale =
            firstMapper.getImageSampleDistance() *
            firstMapper.getImageSampleDistance();
        }
        if (model._lastScale < 1.5) {
          model._lastScale = 1.5;
        }
      });
    }
  };

  publicAPI.rayCastPass = (viewNode, renNode, volumes) => {
    const encoder = model._firstGroup
      ? model._clearEncoder
      : model._mergeEncoder;
    encoder.attachTextureViews();
    encoder.begin(viewNode.getCommandEncoder());
    let width = model._colorTextureView.getTexture().getWidth();
    let height = model._colorTextureView.getTexture().getHeight();
    if (model._useSmallViewport) {
      const canvas = viewNode.getCanvas();
      const scaleFactor = 1 / Math.sqrt(model._lastScale);
      model._smallViewportWidth = Math.ceil(scaleFactor * canvas.width);
      model._smallViewportHeight = Math.ceil(scaleFactor * canvas.height);
      width = model._smallViewportWidth;
      height = model._smallViewportHeight;
    }
    encoder.getHandle().setViewport(0, 0, width, height, 0.0, 1.0);
    // set scissor
    encoder.getHandle().setScissorRect(0, 0, width, height);

    model.fullScreenQuad.setWebGPURenderer(renNode);
    model.fullScreenQuad.setVolumes(volumes);
    model.fullScreenQuad.prepareAndDraw(encoder);
    encoder.end();
  };

  publicAPI.renderDepthBounds = (renNode, viewNode) => {
    publicAPI.updateDepthPolyData(renNode);

    const pd = model._boundsPoly;
    const points = pd.getPoints();
    const cells = pd.getPolys();

    let buffRequest = {
      hash: `vp${cells.getMTime()}`,
      usage: BufferUsage.Index,
      cells,
      numberOfPoints: points.getNumberOfPoints(),
      primitiveType: PrimitiveTypes.Triangles,
      representation: Representation.SURFACE,
    };
    const indexBuffer = viewNode
      .getDevice()
      .getBufferManager()
      .getBuffer(buffRequest);
    model._mapper.getVertexInput().setIndexBuffer(indexBuffer);

    // points
    buffRequest = {
      usage: BufferUsage.PointArray,
      format: 'float32x4',
      hash: `vp${points.getMTime()}${cells.getMTime()}`,
      dataArray: points,
      indexBuffer,
      packExtra: true,
    };
    const buff = viewNode.getDevice().getBufferManager().getBuffer(buffRequest);
    model._mapper.getVertexInput().addBuffer(buff, ['vertexBC']);
    model._mapper.setNumberOfVertices(
      buff.getSizeInBytes() / buff.getStrideInBytes()
    );

    publicAPI.drawDepthRange(renNode, viewNode);
  };

  publicAPI.updateDepthPolyData = (renNode) => {
    // check mtimes first
    let update = false;
    for (let i = 0; i < model.volumes.length; i++) {
      const mtime = model.volumes[i].getMTime();
      if (!model._lastMTimes[i] || mtime !== model._lastMTimes[i]) {
        update = true;
        model._lastMTimes[i] = mtime;
      }
    }

    // also check stabilized time
    const stime = renNode.getStabilizedTime();
    if (
      model._lastMTimes.length <= model.volumes.length ||
      stime !== model._lastMTimes[model.volumes.length]
    ) {
      update = true;
      model._lastMTimes[model.volumes.length] = stime;
    }

    // if no need to update then return
    if (!update) {
      return;
    }

    // rebuild
    const center = renNode.getStabilizedCenterByReference();
    const numPts = model.volumes.length * 8;
    const points = new Float64Array(numPts * 3);
    const numTris = model.volumes.length * 12;
    const polys = new Uint16Array(numTris * 4);

    // add points and cells
    for (let i = 0; i < model.volumes.length; i++) {
      model.volumes[i].getBoundingCubePoints(points, i * 24);
      let cellIdx = i * 12 * 4;
      const offset = i * 8;
      for (let t = 0; t < 12; t++) {
        polys[cellIdx++] = 3;
        polys[cellIdx++] = offset + cubeFaceTriangles[t][0];
        polys[cellIdx++] = offset + cubeFaceTriangles[t][1];
        polys[cellIdx++] = offset + cubeFaceTriangles[t][2];
      }
    }

    for (let p = 0; p < points.length; p += 3) {
      points[p] -= center[0];
      points[p + 1] -= center[1];
      points[p + 2] -= center[2];
    }

    model._boundsPoly.getPoints().setData(points, 3);
    model._boundsPoly.getPoints().modified();
    model._boundsPoly.getPolys().setData(polys, 1);
    model._boundsPoly.getPolys().modified();
    model._boundsPoly.modified();
  };

  publicAPI.drawDepthRange = (renNode, viewNode) => {
    // copy current depth buffer to
    model._depthRangeTexture.resizeToMatch(model.colorTextureView.getTexture());
    model._depthRangeTexture2.resizeToMatch(
      model.colorTextureView.getTexture()
    );

    model._depthRangeEncoder.attachTextureViews();

    publicAPI.setCurrentOperation('volumeDepthRangePass');
    renNode.setRenderEncoder(model._depthRangeEncoder);
    renNode.volumeDepthRangePass(true);
    model._mapper.setWebGPURenderer(renNode);

    model._mapper.prepareToDraw(model._depthRangeEncoder);
    model._mapper.registerDrawCallback(model._depthRangeEncoder);

    renNode.volumeDepthRangePass(false);
  };

  publicAPI.createDepthRangeEncoder = (viewNode) => {
    const device = viewNode.getDevice();
    model._depthRangeEncoder = vtkWebGPURenderEncoder.newInstance({
      label: 'VolumePass DepthRange',
    });
    model._depthRangeEncoder.setPipelineHash('volr');
    model._depthRangeEncoder.setReplaceShaderCodeFunction((pipeline) => {
      const fDesc = pipeline.getShaderDescription('fragment');
      fDesc.addOutput('vec4<f32>', 'outColor1');
      fDesc.addOutput('vec4<f32>', 'outColor2');
      let code = fDesc.getCode();
      code = vtkWebGPUShaderCache.substitute(
        code,
        '//VTK::RenderEncoder::Impl',
        [
          'output.outColor1 = vec4<f32>(input.fragPos.z, 0.0, 0.0, 0.0);',
          'output.outColor2 = vec4<f32>(stopval, 0.0, 0.0, 0.0);',
        ]
      ).result;
      fDesc.setCode(code);
    });
    model._depthRangeEncoder.setDescription({
      colorAttachments: [
        {
          view: null,
          clearValue: [0.0, 0.0, 0.0, 0.0],
          loadOp: 'clear',
          storeOp: 'store',
        },
        {
          view: null,
          clearValue: [1.0, 1.0, 1.0, 1.0],
          loadOp: 'clear',
          storeOp: 'store',
        },
      ],
    });
    model._depthRangeEncoder.setPipelineSettings({
      primitive: { cullMode: 'none' },
      fragment: {
        targets: [
          {
            format: 'r16float',
            blend: {
              color: {
                srcFactor: 'one',
                dstFactor: 'one',
                operation: 'max',
              },
              alpha: { srcfactor: 'one', dstFactor: 'one', operation: 'max' },
            },
          },
          {
            format: 'r16float',
            blend: {
              color: {
                srcFactor: 'one',
                dstFactor: 'one',
                operation: 'min',
              },
              alpha: { srcfactor: 'one', dstFactor: 'one', operation: 'min' },
            },
          },
        ],
      },
    });

    // and the textures it needs
    model._depthRangeTexture = vtkWebGPUTexture.newInstance({
      label: 'volumePassMaxDepth',
    });
    model._depthRangeTexture.create(device, {
      width: viewNode.getCanvas().width,
      height: viewNode.getCanvas().height,
      format: 'r16float',
      usage:
        GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.TEXTURE_BINDING,
    });
    const maxView = model._depthRangeTexture.createView('maxTexture');
    model._depthRangeEncoder.setColorTextureView(0, maxView);
    model._depthRangeTexture2 = vtkWebGPUTexture.newInstance({
      label: 'volumePassDepthMin',
    });
    model._depthRangeTexture2.create(device, {
      width: viewNode.getCanvas().width,
      height: viewNode.getCanvas().height,
      format: 'r16float',
      usage:
        GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.TEXTURE_BINDING,
    });
    const minView = model._depthRangeTexture2.createView('minTexture');
    model._depthRangeEncoder.setColorTextureView(1, minView);
    model._mapper.setDevice(viewNode.getDevice());
    model._mapper.setTextureViews([model.depthTextureView]);
  };

  publicAPI.createClearEncoder = (viewNode) => {
    model._colorTexture = vtkWebGPUTexture.newInstance({
      label: 'volumePassColor',
    });
    model._colorTexture.create(viewNode.getDevice(), {
      width: viewNode.getCanvas().width,
      height: viewNode.getCanvas().height,
      format: 'bgra8unorm',
      /* eslint-disable no-undef */
      /* eslint-disable no-bitwise */
      usage:
        GPUTextureUsage.RENDER_ATTACHMENT |
        GPUTextureUsage.TEXTURE_BINDING |
        GPUTextureUsage.COPY_SRC,
    });
    model._colorTextureView = model._colorTexture.createView(
      'volumePassColorTexture'
    );
    model._colorTextureView.addSampler(viewNode.getDevice(), {
      minFilter: 'linear',
      magFilter: 'linear',
    });

    model._clearEncoder = vtkWebGPURenderEncoder.newInstance({
      label: 'VolumePass Clear',
    });
    model._clearEncoder.setColorTextureView(0, model._colorTextureView);
    model._clearEncoder.setDescription({
      colorAttachments: [
        {
          view: null,
          clearValue: [0.0, 0.0, 0.0, 0.0],
          loadOp: 'clear',
          storeOp: 'store',
        },
      ],
    });
    model._clearEncoder.setPipelineHash('volpf');
    model._clearEncoder.setPipelineSettings({
      primitive: { cullMode: 'none' },
      fragment: {
        targets: [
          {
            format: 'bgra8unorm',
            blend: {
              color: {
                srcFactor: 'src-alpha',
                dstFactor: 'one-minus-src-alpha',
              },
              alpha: { srcfactor: 'one', dstFactor: 'one-minus-src-alpha' },
            },
          },
        ],
      },
    });
  };

  publicAPI.createCopyEncoder = (viewNode) => {
    model._copyEncoder = vtkWebGPURenderEncoder.newInstance({
      label: 'volumePassCopy',
    });
    model._copyEncoder.setDescription({
      colorAttachments: [
        {
          view: null,
          loadOp: 'load',
          storeOp: 'store',
        },
      ],
    });
    model._copyEncoder.setPipelineHash('volcopypf');
    model._copyEncoder.setPipelineSettings({
      primitive: { cullMode: 'none' },
      fragment: {
        targets: [
          {
            format: 'rgba16float',
            blend: {
              color: {
                srcFactor: 'one',
                dstFactor: 'one-minus-src-alpha',
              },
              alpha: { srcfactor: 'one', dstFactor: 'one-minus-src-alpha' },
            },
          },
        ],
      },
    });
  };

  publicAPI.createMergeEncoder = (viewNode) => {
    model._mergeEncoder = vtkWebGPURenderEncoder.newInstance({
      label: 'volumePassMerge',
    });
    model._mergeEncoder.setColorTextureView(0, model._colorTextureView);
    model._mergeEncoder.setDescription({
      colorAttachments: [
        {
          view: null,
          loadOp: 'load',
          storeOp: 'store',
        },
      ],
    });
    model._mergeEncoder.setReplaceShaderCodeFunction((pipeline) => {
      const fDesc = pipeline.getShaderDescription('fragment');
      fDesc.addOutput('vec4<f32>', 'outColor');
      let code = fDesc.getCode();
      code = vtkWebGPUShaderCache.substitute(
        code,
        '//VTK::RenderEncoder::Impl',
        ['output.outColor = vec4<f32>(computedColor.rgb, computedColor.a);']
      ).result;
      fDesc.setCode(code);
    });
    model._mergeEncoder.setPipelineHash('volpf');
    model._mergeEncoder.setPipelineSettings({
      primitive: { cullMode: 'none' },
      fragment: {
        targets: [
          {
            format: 'bgra8unorm',
            blend: {
              color: {
                srcFactor: 'src-alpha',
                dstFactor: 'one-minus-src-alpha',
              },
              alpha: { srcfactor: 'one', dstFactor: 'one-minus-src-alpha' },
            },
          },
        ],
      },
    });
  };

  // marks modified when needed
  publicAPI.setVolumes = (val) => {
    if (!model.volumes || model.volumes.length !== val.length) {
      model.volumes = [...val];
      publicAPI.modified();
      return;
    }
    for (let i = 0; i < val.length; i++) {
      if (val[i] !== model.volumes[i]) {
        model.volumes = [...val];
        publicAPI.modified();
        return;
      }
    }
  };
}

// ----------------------------------------------------------------------------
// Object factory
// ----------------------------------------------------------------------------

const DEFAULT_VALUES = {
  colorTextureView: null,
  depthTextureView: null,
  volumes: null,
};

// ----------------------------------------------------------------------------

export function extend(publicAPI, model, initialValues = {}) {
  Object.assign(model, DEFAULT_VALUES, initialValues);

  // Build VTK API
  vtkRenderPass.extend(publicAPI, model, initialValues);

  model._mapper = vtkWebGPUSimpleMapper.newInstance();
  model._mapper.setFragmentShaderTemplate(DepthBoundsFS);
  model._mapper
    .getShaderReplacements()
    .set('replaceShaderVolumePass', (hash, pipeline, vertexInput) => {
      const fDesc = pipeline.getShaderDescription('fragment');
      fDesc.addBuiltinInput('vec4<f32>', '@builtin(position) fragPos');
    });

  model._boundsPoly = vtkPolyData.newInstance();
  model._lastMTimes = [];

  macro.setGet(publicAPI, model, ['colorTextureView', 'depthTextureView']);

  // Object methods
  vtkWebGPUVolumePass(publicAPI, model);
}

// ----------------------------------------------------------------------------

export const newInstance = macro.newInstance(extend, 'vtkWebGPUVolumePass');

// ----------------------------------------------------------------------------

export default { newInstance, extend };