@babylonjs/viewer/dist/chunks/gs-picking-pipeline-CyS6hN5Z.esm.js

The Babylon Viewer aims to simplify a specific but common Babylon.js use case: loading, viewing, and interacting with a 3D model.

import { B as BU, F as F32, b8 as getRenderTargetSize, b9 as getViewMatrix, ba as getProjectionMatrix, n as SS, aD as getPickingSceneBGL } from './index-DMbDahsc.esm.js';
import { a as applyGsFragments } from './gltf-feature-gaussian-splatting-BiH_nkp6.esm.js';

const gsGpuPickingFragment = {
  id: "gsGpuPicking",
  helperFunctions: (
    /* wgsl */
    `
struct GsPickingU { pickingColor: vec3<f32> };
@group(2) @binding(0) var<uniform> picking: GsPickingU;
`
  ),
  fragmentSlots: {
    GS_FRAGMENT_BEFORE_FRAGCOLOR: (
      /* wgsl */
      `
            if (finalColor.a < 0.001) { discard; }
            finalColor = vec4<f32>(picking.pickingColor, 1.0);
        `
    )
  }
};
function encodeIdToColor(id) {
  return [(id >> 16 & 255) / 255, (id >> 8 & 255) / 255, (id & 255) / 255];
}

let _cache = null;
const _gsPickMx = new F32(16);
function buildPickingWgsl(detailed) {
  const wgsl = (
    /* wgsl */
    `
struct GsPickScene { pickMatrix: mat4x4<f32> };
@group(0) @binding(0) var<uniform> gsPickScene: GsPickScene;

struct U {
  world: mat4x4<f32>,
  view: mat4x4<f32>,
  projection: mat4x4<f32>,
  viewport: vec2<f32>,
  focal: vec2<f32>,
  dataSize: vec2<f32>,
  alpha: f32,
  _pad: f32,
};
@group(1) @binding(0) var<uniform> u: U;
@group(1) @binding(1) var samp: sampler;
@group(1) @binding(2) var centersTex: texture_2d<f32>;
@group(1) @binding(3) var covATex: texture_2d<f32>;
@group(1) @binding(4) var covBTex: texture_2d<f32>;
@group(1) @binding(5) var colorsTex: texture_2d<f32>;

struct VOut {
  @builtin(position) pos: vec4<f32>,
  @location(0) vColor: vec4<f32>,
  @location(1) vPos: vec2<f32>,
};

fn dataUv(idx: f32) -> vec2<f32> {
  let y = floor(idx / u.dataSize.x);
  let x = idx - y * u.dataSize.x;
  return vec2<f32>((x + 0.5) / u.dataSize.x, (y + 0.5) / u.dataSize.y);
}

@vertex
fn vs(@location(0) corner: vec2<f32>, @location(1) splatIndex: f32) -> VOut {
  var out: VOut;
  let uv = dataUv(splatIndex);
  let center = textureSampleLevel(centersTex, samp, uv, 0.0).xyz;
  let color  = textureSampleLevel(colorsTex,  samp, uv, 0.0);
  let covA   = textureSampleLevel(covATex,    samp, uv, 0.0).xyz;
  let covB   = textureSampleLevel(covBTex,    samp, uv, 0.0).xyz;

  let worldPos  = u.world * vec4<f32>(center, 1.0);
  let modelView = u.view  * u.world;
  let camspace  = u.view  * worldPos;
  let pos2d     = u.projection * camspace;

  let bounds = 1.2 * pos2d.w;
  if (pos2d.z < 0.0
      || pos2d.x < -bounds || pos2d.x > bounds
      || pos2d.y < -bounds || pos2d.y > bounds) {
    out.pos = vec4<f32>(0.0, 0.0, 2.0, 1.0);
    out.vColor = vec4<f32>(0.0);
    out.vPos = vec2<f32>(0.0);
    return out;
  }

  let Vrk = mat3x3<f32>(
    vec3<f32>(covA.x, covA.y, covA.z),
    vec3<f32>(covA.y, covB.x, covB.y),
    vec3<f32>(covA.z, covB.y, covB.z));

  let invZ  = 1.0 / camspace.z;
  let invZ2 = invZ * invZ;
  let J = mat3x3<f32>(
    vec3<f32>(u.focal.x * invZ, 0.0, -u.focal.x * camspace.x * invZ2),
    vec3<f32>(0.0, u.focal.y * invZ, -u.focal.y * camspace.y * invZ2),
    vec3<f32>(0.0, 0.0, 0.0));

  let mv3 = mat3x3<f32>(modelView[0].xyz, modelView[1].xyz, modelView[2].xyz);
  let T = transpose(mv3) * J;
  var cov2d = transpose(T) * Vrk * T;

  let kernelSize: f32 = 0.3;
  cov2d[0][0] += kernelSize;
  cov2d[1][1] += kernelSize;

  let mid = (cov2d[0][0] + cov2d[1][1]) * 0.5;
  let dxy = (cov2d[0][0] - cov2d[1][1]) * 0.5;
  let radius = length(vec2<f32>(dxy, cov2d[0][1]));
  let epsilon: f32 = 0.0001;
  let lambda1 = mid + radius + epsilon;
  let lambda2 = mid - radius + epsilon;
  if (lambda2 < 0.0) {
    out.pos = vec4<f32>(0.0, 0.0, 2.0, 1.0);
    out.vColor = vec4<f32>(0.0);
    out.vPos = vec2<f32>(0.0);
    return out;
  }

  let diag = normalize(vec2<f32>(cov2d[0][1], lambda1 - cov2d[0][0]));
  let majorAxis = min(sqrt(2.0 * lambda1), 1024.0) * diag;
  let minorAxis = min(sqrt(2.0 * lambda2), 1024.0) * vec2<f32>(diag.y, -diag.x);

  let vCenter = pos2d.xy;
  out.pos = gsPickScene.pickMatrix * vec4<f32>(
    vCenter + (corner.x * majorAxis + corner.y * minorAxis) * pos2d.w / u.viewport,
    pos2d.z, pos2d.w);
  out.vColor = vec4<f32>(color.rgb, color.a * u.alpha);
  out.vPos = corner;
  return out;
}

/*GS_FRAGMENT_DEFINITIONS*/
struct FsOut { @location(0) color: vec4<f32>, @location(1) depth: f32${detailed ? ", @location(2) detail: vec4u" : ""} };
@fragment
fn fs(in: VOut) -> FsOut {
  /*GS_FRAGMENT_MAIN_BEGIN*/
  let A = -dot(in.vPos, in.vPos);
  var finalColor: vec4<f32>;
  if (A > -4.0) {
    let B = exp(A) * in.vColor.a;
    finalColor = vec4<f32>(in.vColor.rgb, B);
  } else {
    finalColor = vec4<f32>(0.0);
  }
  /*GS_FRAGMENT_BEFORE_FRAGCOLOR*/
  /*GS_FRAGMENT_MAIN_END*/
  return FsOut(finalColor, in.pos.z${detailed ? ", vec4u(0xffffffffu, 0u, 0u, 0u)" : ""});
}
`
  );
  return applyGsFragments(wgsl, [gsGpuPickingFragment]);
}
function getCache(engine) {
  const device = engine._device;
  if (_cache && _cache.device === device) {
    return _cache;
  }
  const meshBGL = device.createBindGroupLayout({
    label: "gs-picking-mesh-bgl",
    entries: [
      { binding: 0, visibility: SS.VERTEX | SS.FRAGMENT, buffer: { type: "uniform" } },
      { binding: 1, visibility: SS.VERTEX, sampler: { type: "non-filtering" } },
      { binding: 2, visibility: SS.VERTEX, texture: { sampleType: "unfilterable-float" } },
      { binding: 3, visibility: SS.VERTEX, texture: { sampleType: "unfilterable-float" } },
      { binding: 4, visibility: SS.VERTEX, texture: { sampleType: "unfilterable-float" } },
      { binding: 5, visibility: SS.VERTEX, texture: { sampleType: "unfilterable-float" } }
    ]
  });
  const pickingBGL = device.createBindGroupLayout({
    label: "gs-picking-pick-bgl",
    entries: [{ binding: 0, visibility: SS.FRAGMENT, buffer: { type: "uniform" } }]
  });
  const pickMatrixUbo = device.createBuffer({ size: 64, usage: BU.UNIFORM | BU.COPY_DST, label: "gs-picking-scene-ubo" });
  const sceneBG = device.createBindGroup({
    label: "gs-picking-scene-bg",
    layout: getPickingSceneBGL(engine),
    entries: [{ binding: 0, resource: { buffer: pickMatrixUbo } }]
  });
  _cache = { device, pipelines: /* @__PURE__ */ new Map(), meshBGL, pickingBGL, pickMatrixUbo, sceneBG };
  return _cache;
}
function getPickingPipeline(engine, detailed) {
  const cache = getCache(engine);
  const key = detailed ? "detailed" : "basic";
  const cached = cache.pipelines.get(key);
  if (cached) {
    return cached;
  }
  const device = engine._device;
  const module = device.createShaderModule({ label: `gs-picking-${key}-shader`, code: buildPickingWgsl(detailed) });
  const pipeline = device.createRenderPipeline({
    label: `gs-picking-${key}-pipeline`,
    layout: device.createPipelineLayout({ bindGroupLayouts: [getPickingSceneBGL(engine), cache.meshBGL, cache.pickingBGL] }),
    vertex: {
      module,
      entryPoint: "vs",
      buffers: [
        { arrayStride: 8, stepMode: "vertex", attributes: [{ shaderLocation: 0, offset: 0, format: "float32x2" }] },
        { arrayStride: 4, stepMode: "instance", attributes: [{ shaderLocation: 1, offset: 0, format: "float32" }] }
      ]
    },
    fragment: {
      module,
      entryPoint: "fs",
      targets: detailed ? [{ format: "rgba8unorm" }, { format: "r32float" }, { format: "rgba32uint" }] : [{ format: "rgba8unorm" }, { format: "r32float" }]
    },
    primitive: { topology: "triangle-list", cullMode: "none" },
    depthStencil: { format: "depth24plus", depthCompare: "less", depthWriteEnabled: true },
    multisample: { count: 1 }
  });
  cache.pipelines.set(key, pipeline);
  return pipeline;
}
function gsPickWritePickMatrixAndBind(pass, engine, pickMatrix) {
  const cache = getCache(engine);
  engine._device.queue.writeBuffer(cache.pickMatrixUbo, 0, pickMatrix.buffer, pickMatrix.byteOffset, pickMatrix.byteLength);
  pass.setBindGroup(0, cache.sceneBG);
}
function createGsPickMeshResources(engine, mesh) {
  const device = engine._device;
  const cache = getCache(engine);
  const UBO_BYTES = 16 * 4 * 3 + 8 * 4;
  const meshUbo = device.createBuffer({ size: UBO_BYTES, usage: BU.UNIFORM | BU.COPY_DST, label: "gs-picking-mesh-ubo" });
  const meshCpu = new F32(UBO_BYTES / 4);
  meshCpu[48 + 4] = mesh.textureWidth;
  meshCpu[48 + 5] = mesh.textureHeight;
  meshCpu[48 + 6] = 1;
  const meshBG = device.createBindGroup({
    label: "gs-picking-mesh-bg",
    layout: cache.meshBGL,
    entries: [
      { binding: 0, resource: { buffer: meshUbo } },
      { binding: 1, resource: mesh._gs._sampler },
      { binding: 2, resource: mesh._gs._centersView },
      { binding: 3, resource: mesh._gs._covAView },
      { binding: 4, resource: mesh._gs._covBView },
      { binding: 5, resource: mesh._gs._colorsView }
    ]
  });
  const pickingUbo = device.createBuffer({ size: 16, usage: BU.UNIFORM | BU.COPY_DST, label: "gs-picking-color-ubo" });
  const pickingCpu = new F32(4);
  const pickingBG = device.createBindGroup({
    label: "gs-picking-color-bg",
    layout: cache.pickingBGL,
    entries: [{ binding: 0, resource: { buffer: pickingUbo } }]
  });
  return { meshUbo, meshBG, pickingUbo, pickingBG, meshCpu, pickingCpu };
}
function disposeGsPickMeshResources(res) {
  res.meshUbo.destroy();
  res.pickingUbo.destroy();
}
function drawGsForPicking(pass, engine, scene, mesh, res, pickId, targetWidth, targetHeight, detailed = false) {
  const cam = scene.camera;
  if (!cam) {
    return;
  }
  const size = getRenderTargetSize(engine);
  const aspect = (targetWidth || size.width) / (targetHeight || size.height);
  const view = getViewMatrix(cam);
  const proj = getProjectionMatrix(cam, aspect);
  const world = mesh.worldMatrix;
  const cpu = res.meshCpu;
  cpu.set(world, 0);
  cpu.set(view, 16);
  cpu.set(proj, 32);
  cpu[48] = size.width;
  cpu[48 + 1] = size.height;
  cpu[48 + 2] = size.width * 0.5 * proj[0];
  cpu[48 + 3] = size.height * 0.5 * proj[5];
  engine._device.queue.writeBuffer(res.meshUbo, 0, cpu.buffer, 0, cpu.byteLength);
  const [r, g, b] = encodeIdToColor(pickId);
  res.pickingCpu[0] = r;
  res.pickingCpu[1] = g;
  res.pickingCpu[2] = b;
  res.pickingCpu[3] = 0;
  engine._device.queue.writeBuffer(res.pickingUbo, 0, res.pickingCpu.buffer, 0, 16);
  pass.setPipeline(getPickingPipeline(engine, detailed));
  pass.setBindGroup(1, res.meshBG);
  pass.setBindGroup(2, res.pickingBG);
  pass.setVertexBuffer(0, mesh._gs._quadBuffer);
  pass.setVertexBuffer(1, mesh._gs._splatIndexBuffer);
  pass.setIndexBuffer(mesh._gs._indexBuffer, "uint16");
  pass.drawIndexed(6, mesh.vertexCount);
}
function drawGsMeshForPicking(ctx, mesh, res, baseId) {
  computeGsPickMatrix(_gsPickMx, ctx.px, ctx.py, ctx.w, ctx.h);
  gsPickWritePickMatrixAndBind(ctx.pass, ctx.engine, _gsPickMx);
  drawGsForPicking(ctx.pass, ctx.engine, ctx.scene, mesh, res, baseId, ctx.w, ctx.h, ctx.detailed);
}
function computeGsPickMatrix(out, sampleX, sampleY, w, h) {
  const ndcX = 2 * sampleX / w - 1;
  const ndcY = 1 - 2 * sampleY / h;
  out[0] = w;
  out[1] = 0;
  out[2] = 0;
  out[3] = 0;
  out[4] = 0;
  out[5] = h;
  out[6] = 0;
  out[7] = 0;
  out[8] = 0;
  out[9] = 0;
  out[10] = 1;
  out[11] = 0;
  out[12] = -ndcX * w;
  out[13] = -ndcY * h;
  out[14] = 0;
  out[15] = 1;
}
function createPickContributor(mesh) {
  let res = null;
  return {
    draw(ctx, baseId) {
      if (mesh.visible === false) {
        return baseId + 1;
      }
      res ??= createGsPickMeshResources(ctx.engine, mesh);
      drawGsMeshForPicking(ctx, mesh, res, baseId);
      return baseId + 1;
    },
    resolve(info) {
      info.pickedMesh = mesh;
    },
    dispose() {
      if (res) {
        disposeGsPickMeshResources(res);
        res = null;
      }
    }
  };
}

export { computeGsPickMatrix, createGsPickMeshResources, createPickContributor, disposeGsPickMeshResources, drawGsForPicking, drawGsMeshForPicking, gsPickWritePickMatrixAndBind };
//# sourceMappingURL=gs-picking-pipeline-CyS6hN5Z.esm.js.map