Add modifiers normalColor and depthColor. Add examples/debug-color to show how to use them on SplatMeshes. (#58)

src/dyno/util.ts

@@ -70,6 +70,13 @@ export const hashVec4 = <T extends ValueTypes>(
 ): DynoVal<"vec4"> => {
   return new HashVec4({ value });
 };
+export const normalizedDepth = (
+  z: DynoVal<"float">,
+  zNear: DynoVal<"float">,
+  zFar: DynoVal<"float">,
+): DynoVal<"float"> => {
+  return new NormalizedDepth({ z, zNear, zFar }).outputs.depth;
+};
 
 export class DynoRemapIndex
   extends Dyno<{ from: "int"; to: "int"; index: "int" }, { index: "int" }>
@@ -404,3 +411,31 @@ export class HashVec4<T extends ValueTypes>
     return new DynoOutput(this, "hash");
   }
 }
+
+export class NormalizedDepth
+  extends Dyno<
+    { z: "float"; zNear: "float"; zFar: "float" },
+    { depth: "float" }
+  >
+  implements HasDynoOut<"float">
+{
+  constructor({
+    z,
+    zNear,
+    zFar,
+  }: { z: DynoVal<"float">; zNear: DynoVal<"float">; zFar: DynoVal<"float"> }) {
+    super({
+      inTypes: { z: "float", zNear: "float", zFar: "float" },
+      outTypes: { depth: "float" },
+      inputs: { z, zNear, zFar },
+      statements: ({ inputs, outputs }) => [
+        `float clamped = clamp(${inputs.z}, ${inputs.zNear}, ${inputs.zFar});`,
+        `${outputs.depth} = (log2(clamped + 1.0) - log2(${inputs.zNear} + 1.0)) / (log2(${inputs.zFar} + 1.0) - log2(${inputs.zNear} + 1.0));`,
+      ],
+    });
+  }
+
+  dynoOut(): DynoValue<"float"> {
+    return new DynoOutput(this, "depth");
+  }
+}

src/index.ts

@@ -48,6 +48,7 @@ export {
 } from "./splatConstructors";
 
 export * as generators from "./generators";
+export * as modifiers from "./modifiers";
 
 export { VRButton } from "./vrButton";
 export {

src/modifiers.ts

@@ -0,0 +1,2 @@
+export * from "./modifiers/normalColor";
+export * from "./modifiers/depthColor";

src/modifiers/depthColor.ts

@@ -0,0 +1,60 @@
+import type { SplatTransformer } from "../SplatGenerator";
+import type { SplatMesh } from "../SplatMesh";
+import {
+  type DynoVal,
+  Gsplat,
+  combineGsplat,
+  dynoBlock,
+  dynoConst,
+  neg,
+  normalizedDepth,
+  select,
+  split,
+  splitGsplat,
+  sub,
+} from "../dyno";
+
+export function makeDepthColorModifier(
+  splatToView: SplatTransformer,
+  minDepth: DynoVal<"float">,
+  maxDepth: DynoVal<"float">,
+  reverse: DynoVal<"bool">,
+) {
+  return dynoBlock({ gsplat: Gsplat }, { gsplat: Gsplat }, ({ gsplat }) => {
+    if (!gsplat) {
+      throw new Error("No gsplat input");
+    }
+    let { center } = splitGsplat(gsplat).outputs;
+    center = splatToView.apply(center);
+    const { z } = split(center).outputs;
+    let depth = normalizedDepth(neg(z), minDepth, maxDepth);
+    depth = select(reverse, sub(dynoConst("float", 1), depth), depth);
+
+    gsplat = combineGsplat({ gsplat, r: depth, g: depth, b: depth });
+    return { gsplat };
+  });
+}
+
+export function setDepthColor(
+  splats: SplatMesh,
+  minDepth: number,
+  maxDepth: number,
+  reverse?: boolean,
+) {
+  splats.enableWorldToView = true;
+  const dynoMinDepth = dynoConst("float", minDepth);
+  const dynoMaxDepth = dynoConst("float", maxDepth);
+  const dynoReverse = dynoConst("bool", reverse ?? false);
+  splats.worldModifier = makeDepthColorModifier(
+    splats.context.worldToView,
+    dynoMinDepth,
+    dynoMaxDepth,
+    dynoReverse,
+  );
+  splats.updateGenerator();
+  return {
+    minDepth: dynoMinDepth,
+    maxDepth: dynoMaxDepth,
+    reverse: dynoReverse,
+  };
+}

src/modifiers/normalColor.ts

@@ -0,0 +1,46 @@
+import type { SplatTransformer } from "../SplatGenerator";
+import type { SplatMesh } from "../SplatMesh";
+import {
+  Gsplat,
+  add,
+  combineGsplat,
+  dot,
+  dynoBlock,
+  dynoConst,
+  greaterThanEqual,
+  gsplatNormal,
+  mul,
+  neg,
+  select,
+  splitGsplat,
+} from "../dyno";
+
+export function makeNormalColorModifier(splatToView: SplatTransformer) {
+  return dynoBlock({ gsplat: Gsplat }, { gsplat: Gsplat }, ({ gsplat }) => {
+    if (!gsplat) {
+      throw new Error("No gsplat input");
+    }
+    let normal = gsplatNormal(gsplat);
+
+    const viewGsplat = splatToView.applyGsplat(gsplat);
+    const viewCenter = splitGsplat(viewGsplat).outputs.center;
+    const viewNormal = gsplatNormal(viewGsplat);
+    const splatDot = dot(viewCenter, viewNormal);
+
+    const sameDir = greaterThanEqual(splatDot, dynoConst("float", 0));
+    normal = select(sameDir, neg(normal), normal);
+    const rgb = add(
+      mul(normal, dynoConst("float", 0.5)),
+      dynoConst("float", 0.5),
+    );
+
+    gsplat = combineGsplat({ gsplat, rgb });
+    return { gsplat };
+  });
+}
+
+export function setWorldNormalColor(splats: SplatMesh) {
+  splats.enableWorldToView = true;
+  splats.worldModifier = makeNormalColorModifier(splats.context.worldToView);
+  splats.updateGenerator();
+}