@takram/three-geospatial 0.3.0 → 0.5.0

package/src/webgpu/MipmapSurfaceBlurNode.ts

@@ -0,0 +1,140 @@
+import { add, Fn, mix, nodeObject, uniform, uv, vec2, vec4 } from 'three/tsl'
+import type { NodeBuilder, TextureNode } from 'three/webgpu'
+import invariant from 'tiny-invariant'
+
+import { DualMipmapFilterNode } from './DualMipmapFilterNode'
+import type { Node, NodeObject } from './node'
+
+const clampToBorder = (uv: NodeObject<'vec2'>): NodeObject<'float'> => {
+  return uv.greaterThanEqual(0).all().and(uv.lessThanEqual(1).all()).toFloat()
+}
+
+// Implementation of Lena Piquet's bloom filter.
+// Reference: https://www.froyok.fr/blog/2021-12-ue4-custom-bloom
+export class MipmapSurfaceBlurNode extends DualMipmapFilterNode {
+  static override get type(): string {
+    return 'MipmapSurfaceBlurNode'
+  }
+
+  blendAmount = uniform(0.85)
+
+  constructor(inputNode?: TextureNode | null, levels = 4) {
+    super(inputNode, levels)
+    this.resolutionScale = 0.5
+  }
+
+  protected override setupDownsampleNode(builder: NodeBuilder): Node {
+    const { inputNode, inputTexelSize } = this
+    invariant(inputNode != null)
+
+    return Fn(() => {
+      const center = uv()
+      const offset1 = vec4(1, 1, -1, -1)
+        .mul(inputTexelSize.xyxy)
+        .add(center.xyxy)
+      const offset2 = vec4(2, 2, -2, -2)
+        .mul(inputTexelSize.xyxy)
+        .add(center.xyxy)
+      const uv01 = offset1.zy.toVertexStage() // -1, 1
+      const uv02 = offset1.xy.toVertexStage() // 1, 1
+      const uv03 = offset1.zw.toVertexStage() // -1, -1
+      const uv04 = offset1.xw.toVertexStage() // 1, -1
+      const uv05 = offset2.zy.toVertexStage() // -2, 2
+      const uv06 = offset2.xy.toVertexStage() // 2, 2
+      const uv07 = offset2.zw.toVertexStage() // -2, -2
+      const uv08 = offset2.xw.toVertexStage() // 2, -2
+      const uv09 = vec2(center.x, offset2.y).toVertexStage() // 0, 2
+      const uv10 = vec2(offset2.z, center.y).toVertexStage() // -2, 0
+      const uv11 = vec2(offset2.x, center.y).toVertexStage() // 2, 0
+      const uv12 = vec2(center.x, offset2.w).toVertexStage() // 0, -2
+
+      const innerWeight = 1 / 4 / 2
+      const outerWeight = 1 / 9 / 2
+
+      const output = inputNode.sample(center).mul(outerWeight)
+
+      let weight: NodeObject
+      weight = vec4(
+        clampToBorder(uv01),
+        clampToBorder(uv02),
+        clampToBorder(uv03),
+        clampToBorder(uv04)
+      ).mul(innerWeight)
+
+      output.addAssign(
+        inputNode.sample(uv01).mul(weight.x),
+        inputNode.sample(uv02).mul(weight.y),
+        inputNode.sample(uv03).mul(weight.z),
+        inputNode.sample(uv04).mul(weight.w)
+      )
+
+      weight = vec4(
+        clampToBorder(uv05),
+        clampToBorder(uv06),
+        clampToBorder(uv07),
+        clampToBorder(uv08)
+      ).mul(outerWeight)
+
+      output.addAssign(
+        inputNode.sample(uv05).mul(weight.x),
+        inputNode.sample(uv06).mul(weight.y),
+        inputNode.sample(uv07).mul(weight.z),
+        inputNode.sample(uv08).mul(weight.w)
+      )
+
+      weight = vec4(
+        clampToBorder(uv09),
+        clampToBorder(uv10),
+        clampToBorder(uv11),
+        clampToBorder(uv12)
+      ).mul(outerWeight)
+
+      output.addAssign(
+        inputNode.sample(uv09).mul(weight.x),
+        inputNode.sample(uv10).mul(weight.y),
+        inputNode.sample(uv11).mul(weight.z),
+        inputNode.sample(uv12).mul(weight.w)
+      )
+
+      return output
+    })()
+  }
+
+  protected override setupUpsampleNode(builder: NodeBuilder): Node {
+    const { inputNode, inputTexelSize, downsampleNode } = this
+    invariant(inputNode != null)
+
+    const center = uv()
+    const offset = vec4(1, 1, -1, -1).mul(inputTexelSize.xyxy).add(center.xyxy)
+    const uv1 = vec2(center.x, offset.y).toVertexStage() // 0, 1
+    const uv2 = vec2(offset.z, center.y).toVertexStage() // -1, 0
+    const uv3 = vec2(offset.x, center.y).toVertexStage() // 1, 0
+    const uv4 = vec2(center.x, offset.w).toVertexStage() // 0, -1
+    const uv5 = offset.zy.toVertexStage() // -1, 1
+    const uv6 = offset.xy.toVertexStage() // 1, 1
+    const uv7 = offset.zw.toVertexStage() // -1, -1
+    const uv8 = offset.xw.toVertexStage() // 1, -1
+
+    const output = add(
+      inputNode.sample(center).mul(1 / 4),
+      add(
+        inputNode.sample(uv1),
+        inputNode.sample(uv2),
+        inputNode.sample(uv3),
+        inputNode.sample(uv4)
+      ).mul(1 / 8),
+      add(
+        inputNode.sample(uv5),
+        inputNode.sample(uv6),
+        inputNode.sample(uv7),
+        inputNode.sample(uv8)
+      ).mul(1 / 16)
+    )
+    return mix(downsampleNode.sample(center), output, this.blendAmount)
+  }
+}
+
+export const mipmapSurfaceBlur = (
+  ...args: ConstructorParameters<typeof MipmapSurfaceBlurNode>
+): NodeObject<MipmapSurfaceBlurNode> =>
+  nodeObject(new MipmapSurfaceBlurNode(...args))

package/src/webgpu/OutputTexture3DNode.ts

@@ -0,0 +1,34 @@
+import type { Texture } from 'three'
+import { nodeObject } from 'three/tsl'
+import { Texture3DNode, type Node, type NodeBuilder } from 'three/webgpu'
+
+import type { NodeObject } from './node'
+
+export class OutputTexture3DNode extends Texture3DNode {
+  static override get type(): string {
+    return 'OutputTexture3DNode'
+  }
+
+  owner: Node
+
+  constructor(owner: Node, texture: Texture) {
+    super(texture)
+    this.owner = owner
+    this.setUpdateMatrix(false)
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    this.owner.build(builder)
+    return super.setup(builder)
+  }
+
+  override clone(): this {
+    // @ts-expect-error Ignore
+    return new this.constructor(this.owner, this.value)
+  }
+}
+
+export const outputTexture3D = (
+  ...args: ConstructorParameters<typeof OutputTexture3DNode>
+): NodeObject<OutputTexture3DNode> =>
+  nodeObject(new OutputTexture3DNode(...args))

package/src/webgpu/OutputTextureNode.ts

@@ -0,0 +1,33 @@
+import type { Texture } from 'three'
+import { nodeObject } from 'three/tsl'
+import { TextureNode, type Node, type NodeBuilder } from 'three/webgpu'
+
+import type { NodeObject } from './node'
+
+export class OutputTextureNode extends TextureNode {
+  static override get type(): string {
+    return 'OutputTextureNode'
+  }
+
+  owner: Node
+
+  constructor(owner: Node, texture: Texture) {
+    super(texture)
+    this.owner = owner
+    this.setUpdateMatrix(false)
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    this.owner.build(builder)
+    return super.setup(builder)
+  }
+
+  override clone(): this {
+    // @ts-expect-error Ignore
+    return new this.constructor(this.owner, this.value)
+  }
+}
+
+export const outputTexture = (
+  ...args: ConstructorParameters<typeof OutputTextureNode>
+): NodeObject<OutputTextureNode> => nodeObject(new OutputTextureNode(...args))

package/src/webgpu/RTTextureNode.ts

@@ -0,0 +1,132 @@
+import { nodeObject, uv } from 'three/tsl'
+import {
+  HalfFloatType,
+  LinearFilter,
+  NodeMaterial,
+  NodeUpdateType,
+  QuadMesh,
+  RendererUtils,
+  RenderTarget,
+  RGBAFormat,
+  TextureNode,
+  Vector2,
+  type Node,
+  type NodeBuilder,
+  type NodeFrame
+} from 'three/webgpu'
+
+import type { NodeObject } from './node'
+
+const { resetRendererState, restoreRendererState } = RendererUtils
+
+function createRenderTarget(): RenderTarget {
+  const renderTarget = new RenderTarget(1, 1, {
+    depthBuffer: false,
+    type: HalfFloatType,
+    format: RGBAFormat
+  })
+  const texture = renderTarget.texture
+  texture.minFilter = LinearFilter
+  texture.magFilter = LinearFilter
+  texture.generateMipmaps = false
+  return renderTarget
+}
+
+const sizeScratch = /*#__PURE__*/ new Vector2()
+
+// Similar to RTTNode, which is a bit finicky to handle.
+export class RTTextureNode extends TextureNode {
+  static override get type(): string {
+    return 'RTTextureNode'
+  }
+
+  node: Node
+  resolutionScale = 1
+
+  private readonly renderTarget: RenderTarget
+  private readonly material = new NodeMaterial()
+  private readonly mesh = new QuadMesh(this.material)
+  private rendererState?: RendererUtils.RendererState
+
+  constructor(node: Node, uvNode?: Node) {
+    const renderTarget = createRenderTarget()
+    super(renderTarget.texture, uvNode != null ? nodeObject(uvNode) : uv())
+    this.node = node
+    this.renderTarget = renderTarget
+    this.updateBeforeType = NodeUpdateType.FRAME
+  }
+
+  setSize(width: number, height: number): this {
+    const { resolutionScale } = this
+    const w = Math.max(Math.round(width * resolutionScale), 1)
+    const h = Math.max(Math.round(height * resolutionScale), 1)
+    this.renderTarget.setSize(w, h)
+    return this
+  }
+
+  override updateBefore({ renderer }: NodeFrame): void {
+    if (renderer == null) {
+      return
+    }
+
+    const size = renderer.getDrawingBufferSize(sizeScratch)
+    this.setSize(size.x, size.y)
+
+    this.rendererState = resetRendererState(renderer, this.rendererState)
+
+    renderer.setRenderTarget(this.renderTarget)
+    this.mesh.render(renderer)
+
+    restoreRendererState(renderer, this.rendererState)
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    const { material } = this
+    // I don't fully understand why, but updates on "node" doesn't propagate
+    // unless giving the builder context.
+    material.fragmentNode = nodeObject(this.node).context(builder.getContext())
+    material.needsUpdate = true
+    return super.setup(builder)
+  }
+
+  override dispose(): void {
+    this.renderTarget.dispose()
+    this.material.dispose()
+    this.mesh.geometry.dispose()
+    super.dispose()
+  }
+
+  // @ts-expect-error Ignore
+  override clone(): TextureNode {
+    const result = new TextureNode(this.value, this.uvNode, this.levelNode)
+    result.sampler = this.sampler
+    result.referenceNode = this
+    return result
+  }
+}
+
+export const rtTexture = (
+  ...args: ConstructorParameters<typeof RTTextureNode>
+): NodeObject<RTTextureNode> => nodeObject(new RTTextureNode(...args))
+
+export const convertToTexture = (
+  node: Node & {
+    isTextureNode?: boolean
+    isSampleNode?: boolean
+    getTextureNode?: () => TextureNode
+  },
+  name?: string
+): NodeObject<TextureNode> => {
+  let textureNode: TextureNode
+  if (node.isTextureNode === true || node.isSampleNode === true) {
+    textureNode = node as TextureNode
+  } else if (node.getTextureNode != null) {
+    textureNode = node.getTextureNode()
+  } else {
+    textureNode = new RTTextureNode(node)
+    if (name != null) {
+      textureNode.value.name = name
+    }
+  }
+  return nodeObject(textureNode)
+}

package/src/webgpu/SeparableFilterNode.ts

@@ -0,0 +1,98 @@
+import { Vector2, type RenderTarget } from 'three'
+import { uniform } from 'three/tsl'
+import {
+  NodeMaterial,
+  QuadMesh,
+  RendererUtils,
+  type NodeBuilder,
+  type NodeFrame,
+  type TextureNode
+} from 'three/webgpu'
+import invariant from 'tiny-invariant'
+
+import { FilterNode } from './FilterNode'
+import type { Node } from './node'
+
+const { resetRendererState, restoreRendererState } = RendererUtils
+
+export abstract class SeparableFilterNode extends FilterNode {
+  iterations = 1
+
+  private readonly horizontalRT: RenderTarget
+  private readonly verticalRT: RenderTarget
+  private readonly material = new NodeMaterial()
+  private readonly mesh = new QuadMesh(this.material)
+  private rendererState?: RendererUtils.RendererState
+
+  protected readonly inputTexelSize = uniform(new Vector2())
+  protected readonly direction = uniform(new Vector2())
+
+  constructor(inputNode?: TextureNode | null) {
+    super(inputNode)
+
+    this.horizontalRT = this.createRenderTarget('Horizontal')
+    this.verticalRT = this.createRenderTarget('Vertical')
+    this.outputTexture = this.verticalRT.texture
+  }
+
+  setSize(width: number, height: number): this {
+    const { resolutionScale } = this
+    const w = Math.max(Math.round(width * resolutionScale), 1)
+    const h = Math.max(Math.round(height * resolutionScale), 1)
+    this.horizontalRT.setSize(w, h)
+    this.verticalRT.setSize(w, h)
+    return this
+  }
+
+  override updateBefore({ renderer }: NodeFrame): void {
+    if (renderer == null) {
+      return
+    }
+
+    const { horizontalRT, verticalRT, mesh, inputNode, direction } = this
+    invariant(inputNode != null)
+
+    const { width, height } = inputNode.value
+    this.setSize(width, height)
+    this.inputTexelSize.value.set(1 / width, 1 / height)
+
+    const originalTexture = inputNode.value
+    this.rendererState = resetRendererState(renderer, this.rendererState)
+
+    for (let i = 0; i < this.iterations; ++i) {
+      direction.value.set(1, 0)
+      renderer.setRenderTarget(horizontalRT)
+      mesh.render(renderer)
+      inputNode.value = horizontalRT.texture
+
+      direction.value.set(0, 1)
+      renderer.setRenderTarget(verticalRT)
+      mesh.render(renderer)
+      inputNode.value = verticalRT.texture
+    }
+
+    restoreRendererState(renderer, this.rendererState)
+    inputNode.value = originalTexture
+  }
+
+  protected abstract setupOutputNode(builder: NodeBuilder): Node
+
+  override setup(builder: NodeBuilder): unknown {
+    const { inputNode } = this
+    invariant(inputNode != null)
+
+    const { material } = this
+    material.fragmentNode = this.setupOutputNode(builder)
+    material.needsUpdate = true
+
+    return super.setup(builder)
+  }
+
+  override dispose(): void {
+    this.horizontalRT.dispose()
+    this.verticalRT.dispose()
+    this.material.dispose()
+    this.mesh.geometry.dispose()
+    super.dispose()
+  }
+}

package/src/webgpu/SingleFilterNode.ts

@@ -0,0 +1,80 @@
+import { Vector2, type RenderTarget } from 'three'
+import { uniform } from 'three/tsl'
+import {
+  NodeMaterial,
+  QuadMesh,
+  RendererUtils,
+  type NodeBuilder,
+  type NodeFrame,
+  type TextureNode
+} from 'three/webgpu'
+import invariant from 'tiny-invariant'
+
+import { FilterNode } from './FilterNode'
+import type { Node } from './node'
+
+const { resetRendererState, restoreRendererState } = RendererUtils
+
+export abstract class SingleFilterNode extends FilterNode {
+  private readonly renderTarget: RenderTarget
+  private readonly material = new NodeMaterial()
+  private readonly mesh = new QuadMesh(this.material)
+  private rendererState?: RendererUtils.RendererState
+
+  protected readonly inputTexelSize = uniform(new Vector2())
+
+  constructor(inputNode?: TextureNode | null) {
+    super(inputNode)
+
+    this.renderTarget = this.createRenderTarget()
+    this.outputTexture = this.renderTarget.texture
+  }
+
+  override setSize(width: number, height: number): this {
+    const { resolutionScale } = this
+    const w = Math.max(Math.round(width * resolutionScale), 1)
+    const h = Math.max(Math.round(height * resolutionScale), 1)
+    this.renderTarget.setSize(w, h)
+    return this
+  }
+
+  override updateBefore({ renderer }: NodeFrame): void {
+    if (renderer == null) {
+      return
+    }
+
+    const { inputNode } = this
+    invariant(inputNode != null)
+
+    const { width, height } = inputNode.value
+    this.setSize(width, height)
+    this.inputTexelSize.value.set(1 / width, 1 / height)
+
+    this.rendererState = resetRendererState(renderer, this.rendererState)
+
+    renderer.setRenderTarget(this.renderTarget)
+    this.mesh.render(renderer)
+
+    restoreRendererState(renderer, this.rendererState)
+  }
+
+  protected abstract setupOutputNode(builder: NodeBuilder): Node
+
+  override setup(builder: NodeBuilder): unknown {
+    const { inputNode } = this
+    invariant(inputNode != null)
+
+    const { material } = this
+    material.fragmentNode = this.setupOutputNode(builder)
+    material.needsUpdate = true
+
+    return super.setup(builder)
+  }
+
+  override dispose(): void {
+    this.renderTarget.dispose()
+    this.material.dispose()
+    this.mesh.geometry.dispose()
+    super.dispose()
+  }
+}