feat(web): support HBAO (#569)

reearth · Jul 12, 2023 · d5b3c94 · d5b3c94
1 parent 641a310
commit d5b3c94
Show file tree

Hide file tree

Showing 24 changed files with 1,154 additions and 0 deletions.
diff --git a/web/src/beta/lib/core/Map/types/index.ts b/web/src/beta/lib/core/Map/types/index.ts
@@ -318,6 +318,12 @@ export type SceneProperty = {
     themeSelectColor?: string;
     themeBackgroundColor?: string;
   };
+  ambientOcclusion?: {
+    enabled?: boolean;
+    quality?: "low" | "medium" | "high" | "extreme";
+    intensity?: number;
+    ambientOcclusionOnly?: boolean;
+  };
   light?: {
     lightType?: "sunLight" | "directionalLight";
     lightDirectionX?: number;

diff --git a/web/src/beta/lib/core/engines/Cesium/PostProcesses/README.md b/web/src/beta/lib/core/engines/Cesium/PostProcesses/README.md
@@ -0,0 +1 @@
+# Post processes
diff --git a/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/AmbientOcclusion.tsx b/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/AmbientOcclusion.tsx
@@ -0,0 +1,103 @@
+import { PerspectiveFrustum, type Scene } from "@cesium/engine";
+import { useState, type FC } from "react";
+import { useCesium } from "resium";
+import invariant from "tiny-invariant";
+
+import { useInstance } from "../../hooks/useInstance";
+import { usePreRender } from "../../hooks/useSceneEvent";
+
+import {
+  createAmbientOcclusionStage,
+  type AmbientOcclusionStageOptions,
+  type AmbientOcclusionStageUniforms,
+} from "./createAmbientOcclusionStage";
+
+export interface AmbientOcclusionProps
+  extends Omit<AmbientOcclusionStageOptions, "prefix">,
+    Partial<AmbientOcclusionStageUniforms> {
+  enabled?: boolean;
+}
+
+export const AmbientOcclusionStage = ({
+  enabled = true,
+  textureScale,
+  directions,
+  steps,
+  denoise,
+  accurateNormalReconstruction,
+  outputType,
+  useGlobeDepth,
+  ...uniforms
+}: AmbientOcclusionProps) => {
+  const { viewer } = useCesium();
+  const scene = viewer?.scene;
+  const stage = useInstance({
+    owner: scene?.postProcessStages,
+    keys: [
+      textureScale,
+      steps,
+      directions,
+      denoise,
+      accurateNormalReconstruction,
+      outputType,
+      useGlobeDepth,
+    ],
+    create: () => {
+      const stage = createAmbientOcclusionStage({
+        textureScale,
+        steps,
+        directions,
+        denoise,
+        accurateNormalReconstruction,
+        outputType,
+        useGlobeDepth,
+        // It's tricky to bind globe's depth texture to postprocess stage,
+        // because it is undefined for the first several frames, or when
+        // terrain is being reconstructed. Functional uniform value is invoked
+        // when this stage is actually being rendered.
+        getGlobeDepthTexture: () =>
+          (
+            scene as Scene & {
+              context: {
+                uniformState: {
+                  globeDepthTexture?: unknown;
+                };
+              };
+            }
+          ).context.uniformState.globeDepthTexture,
+        uniforms,
+      });
+      stage.enabled = enabled;
+      return stage;
+    },
+    transferOwnership: (stage, postProcessStages) => {
+      postProcessStages?.add(stage);
+      return () => {
+        postProcessStages?.remove(stage);
+      };
+    },
+  });
+
+  stage.enabled = enabled;
+  Object.assign(stage.uniforms, uniforms);
+  scene?.requestRender();
+
+  usePreRender(() => {
+    const frustum = scene?.camera.frustum;
+    invariant(frustum instanceof PerspectiveFrustum);
+    const cotFovy = 1 / Math.tan(frustum.fovy / 2);
+    stage.uniforms.focalLength.x = cotFovy * frustum.aspectRatio;
+    stage.uniforms.focalLength.y = cotFovy;
+  });
+
+  return null;
+};
+
+export const AmbientOcclusion: FC<AmbientOcclusionProps> = props => {
+  const viewer = useCesium();
+  const [useGlobeDepth, setUseGlobeDepth] = useState(false);
+  usePreRender(() => {
+    setUseGlobeDepth(!viewer.scene?.globe.depthTestAgainstTerrain);
+  });
+  return <AmbientOcclusionStage {...props} useGlobeDepth={useGlobeDepth} />;
+};
diff --git a/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/AmbientOcclusionOutputType.ts b/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/AmbientOcclusionOutputType.ts
@@ -0,0 +1,7 @@
+export const enum AmbientOcclusionOutputType {
+  Occlusion = 1,
+  Normal = 2,
+  Depth = 3,
+  Weight = 4,
+  Shade = 5,
+}
diff --git a/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/README.md b/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/README.md
@@ -0,0 +1,3 @@
+# HBAO post-processor
+
+An implementation forked from [takram-design-enginering/plateau-view](https://github.com/takram-design-engineering/plateau-view/tree/main/libs/cesium-hbao).
diff --git a/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/config.ts b/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/config.ts
@@ -0,0 +1,29 @@
+export const AMBIENT_OCCLUSION_QUALITY: Record<
+  "low" | "medium" | "high" | "extreme",
+  { steps: number; directions: number; textureScale: number; maxRadius: number }
+> = {
+  low: {
+    directions: 4,
+    steps: 4,
+    textureScale: 0.5,
+    maxRadius: 30,
+  },
+  medium: {
+    directions: 4,
+    steps: 8,
+    textureScale: 0.5,
+    maxRadius: 40,
+  },
+  high: {
+    directions: 8,
+    steps: 8,
+    textureScale: 1,
+    maxRadius: 40,
+  },
+  extreme: {
+    directions: 16,
+    steps: 16,
+    textureScale: 1,
+    maxRadius: 40,
+  },
+};
diff --git a/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/createAmbientOcclusionStage.ts b/web/src/beta/lib/core/engines/Cesium/PostProcesses/hbao/createAmbientOcclusionStage.ts
@@ -0,0 +1,173 @@
+import { Cartesian2, Color, PostProcessStage, PostProcessStageComposite } from "@cesium/engine";
+import { defaults, pick } from "lodash-es";
+
+import { createUniforms } from "../../helpers/createUniforms";
+
+import { type AmbientOcclusionOutputType } from "./AmbientOcclusionOutputType";
+import { createBilateralFilterStage } from "./createCrossBilateralFilterStage";
+import ambientOcclusionGenerate from "./shaders/ambientOcclusionGenerate.glsl?raw";
+import ambientOcclusionModulate from "./shaders/ambientOcclusionModulate.glsl?raw";
+import depth from "./shaders/depth.glsl?raw";
+import globeDepth from "./shaders/globeDepth.glsl?raw";
+import highPassRandom from "./shaders/highPassRandom.glsl?raw";
+import packing from "./shaders/packing.glsl?raw";
+import reconstructNormal from "./shaders/reconstructNormal.glsl?raw";
+import reconstructPosition from "./shaders/reconstructPosition.glsl?raw";
+import turboColorMap from "./shaders/turboColorMap.glsl?raw";
+
+interface PrivatePostProcessStage extends PostProcessStage {
+  _depthTexture: unknown;
+}
+
+export interface AmbientOcclusionStageUniforms {
+  intensity: number;
+  color: Color;
+  maxRadius: number;
+  frustumLength: number;
+  bias: number;
+  focalLength: Cartesian2;
+  blackPoint: number;
+  whitePoint: number;
+  gamma: number;
+  normalExponent?: number;
+  depthExponent?: number;
+}
+
+const defaultUniforms: Omit<AmbientOcclusionStageUniforms, "globeDepthTexture"> = {
+  intensity: 100,
+  color: Color.BLACK,
+  maxRadius: 30,
+  bias: 0.1,
+  frustumLength: 1e5,
+  focalLength: new Cartesian2(),
+  blackPoint: 0.05,
+  whitePoint: 0.9,
+  gamma: 2.5,
+};
+
+export interface AmbientOcclusionStageOptions {
+  prefix?: string;
+  textureScale?: number;
+  directions?: number;
+  steps?: number;
+  denoise?: boolean;
+  accurateNormalReconstruction?: boolean;
+  outputType?: AmbientOcclusionOutputType | null;
+  useGlobeDepth?: boolean;
+  getGlobeDepthTexture?: () => unknown | undefined;
+  uniforms?: Partial<AmbientOcclusionStageUniforms>;
+}
+
+export function createAmbientOcclusionStage({
+  prefix = "reearth",
+  textureScale = 1,
+  directions,
+  steps,
+  denoise = true,
+  accurateNormalReconstruction = true,
+  outputType = null,
+  useGlobeDepth = false,
+  getGlobeDepthTexture,
+  uniforms: uniformsOption = {},
+}: AmbientOcclusionStageOptions): PostProcessStageComposite {
+  const uniforms = defaults({}, uniformsOption, defaultUniforms);
+
+  const generate = new PostProcessStage({
+    name: `${prefix}_ambient_occlusion_generate`,
+    fragmentShader: `
+      ${directions != null ? `#define NUM_DIRECTIONS (${directions})` : ""}
+      ${steps != null ? `#define NUM_STEPS (${steps})` : ""}
+      ${accurateNormalReconstruction ? "#define USE_ACCURATE_NORMAL" : ""}
+      ${outputType != null ? `#define OUTPUT_TYPE (${outputType})` : ""}
+      ${useGlobeDepth ? globeDepth : depth}
+      ${packing}
+      ${highPassRandom}
+      ${reconstructPosition}
+      ${reconstructNormal}
+      ${ambientOcclusionGenerate}
+    `,
+    textureScale,
+    uniforms: {
+      textureScale,
+      ...pick(uniforms, ["intensity", "maxRadius", "bias", "frustumLength", "focalLength"]),
+      globeDepthTexture: () => {
+        return getGlobeDepthTexture?.() ?? (generate as PrivatePostProcessStage)._depthTexture;
+      },
+    },
+  });
+
+  const modulate = new PostProcessStage({
+    name: `${prefix}_ambient_occlusion_modulate`,
+    fragmentShader: `
+      ${outputType != null ? `#define OUTPUT_TYPE (${outputType})` : ""}
+      ${useGlobeDepth ? globeDepth : depth}
+      ${packing}
+      ${reconstructPosition}
+      ${turboColorMap}
+      ${ambientOcclusionModulate}
+    `,
+    uniforms: {
+      ...pick(uniforms, ["color", "frustumLength", "blackPoint", "whitePoint", "gamma"]),
+      globeDepthTexture: () =>
+        getGlobeDepthTexture?.() ?? (modulate as PrivatePostProcessStage)._depthTexture,
+    },
+  });
+
+  if (denoise) {
+    const filter = createBilateralFilterStage({
+      prefix,
+      outputType,
+      useGlobeDepth,
+      getGlobeDepthTexture,
+      uniforms: {
+        textureScale,
+        ...pick(uniforms, ["frustumLength", "normalExponent", "depthExponent"]),
+      },
+    });
+
+    const generateAndFilter = new PostProcessStageComposite({
+      name: `${prefix}_ambient_occlusion_generate_and_filter`,
+      stages: [generate, filter],
+    });
+
+    modulate.uniforms.ambientOcclusionTexture = generateAndFilter.name;
+
+    return new PostProcessStageComposite({
+      name: `${prefix}_ambient_occlusion`,
+      stages: [generateAndFilter, modulate],
+      inputPreviousStageTexture: false,
+      uniforms: createUniforms<AmbientOcclusionStageUniforms>([
+        {
+          stage: generate,
+          uniforms: ["intensity", "maxRadius", "bias", "frustumLength", "focalLength"],
+        },
+        {
+          stage: filter,
+          uniforms: ["frustumLength", "normalExponent", "depthExponent"],
+        },
+        {
+          stage: modulate,
+          uniforms: ["color", "blackPoint", "whitePoint", "gamma"],
+        },
+      ]),
+    });
+  } else {
+    modulate.uniforms.ambientOcclusionTexture = generate.name;
+
+    return new PostProcessStageComposite({
+      name: `${prefix}_ambient_occlusion`,
+      stages: [generate, modulate],
+      inputPreviousStageTexture: false,
+      uniforms: createUniforms<AmbientOcclusionStageUniforms>([
+        {
+          stage: generate,
+          uniforms: ["intensity", "maxRadius", "bias", "frustumLength", "focalLength"],
+        },
+        {
+          stage: modulate,
+          uniforms: ["color", "blackPoint", "whitePoint", "gamma"],
+        },
+      ]),
+    });
+  }
+}