Merge pull request #7 from structureio/feature/new-depthOverlay

kazi-mifta · web-flow · commit 42c2cea9018f · 2025-09-19T17:48:44.000-06:00
[ST-2164] (+) Depth overlay range based
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ## Unreleased
 
+
+## 1.0.6
+
+### Added
+* Range based depth overlay rendering. Depth will be colored differently based on inside and outside of a specified range.
+
 ## 1.0.5
 
 ### Added
diff --git a/Sources/StructureKit/Metal/STKDepthBandOverlay.metal b/Sources/StructureKit/Metal/STKDepthBandOverlay.metal
@@ -0,0 +1,86 @@
+// StructureKit - A collection of extension utilities for Structure SDK
+// Copyright 2022 XRPro, LLC. All rights reserved.
+// http://structure.io
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+// * Neither the name of XRPro, LLC nor the names of its contributors may be
+//   used to endorse or promote products derived from this software without
+//   specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+
+#include "STKMetalCommon.h"
+#include "STKMetalData.h"
+
+#include <metal_stdlib>
+
+using namespace metal;
+
+vertex STKVertexTexOut vertexDepthBandOverlay(
+    const device STKVertexTex* vertex_array [[buffer(0)]],
+    const device STKUniformsDepthOverlay& uniforms [[buffer(1)]],
+    unsigned int vid [[vertex_id]])
+{
+    const STKVertexTex VertexIn = vertex_array[vid];
+    const float4 texTrans = uniforms.projection * float4(VertexIn.texCoord, 0, 1);
+
+    STKVertexTexOut VertexOut;
+    VertexOut.position = float4(VertexIn.position, 1);
+    VertexOut.texCoord = float2(texTrans.x, texTrans.y);
+    return VertexOut;
+}
+
+fragment float4 fragmentDepthBandOverlay(
+    STKVertexTexOut interpolated [[stage_in]],
+    const device STKUniformsDepthBandOverlay& uniforms [[buffer(0)]],
+    texture2d<float> texDepth [[texture(0)]],
+    sampler sampler2D [[sampler(0)]])
+{
+    // depth udefined
+    const float depth = texDepth.sample(sampler2D, interpolated.texCoord).r;
+    if (isnan(depth))
+        return float4(0);
+
+    // calculate position of the depth pixel in the world CS using intrinsics
+    const auto intrinsics = uniforms.cameraIntrinsics;
+    const float u = interpolated.texCoord.x * intrinsics.width;
+    const float v = interpolated.texCoord.y * intrinsics.height;
+    const float depthM = depth / 1000;
+    const float4 cameraPoint(
+        depthM * (u - intrinsics.cx) / intrinsics.fx,
+        depthM * (v - intrinsics.cy) / intrinsics.fy,
+        depthM,
+        1);
+    const float4 worldPoint = uniforms.cameraPose * cameraPoint;
+    const float4 cubePoint = uniforms.cubeModelInv * worldPoint;
+
+    // outside the box
+    if (cubePoint.x < 0 || cubePoint.x > 1 || cubePoint.y < 0 || cubePoint.y > 1 || cubePoint.z < 0 || cubePoint.z > 1)
+        return float4(0);
+
+    // Use validRangeColor in ideal range, outOfRangeColor elsewhere.
+    float inAtMin  = smoothstep(uniforms.validRangeMinMM - uniforms.feather, uniforms.validRangeMinMM + uniforms.feather, depth);
+    float inAtMax  = 1.0 - smoothstep(uniforms.validRangeMaxMM - uniforms.feather, uniforms.validRangeMaxMM + uniforms.feather, depth);
+    float t        = clamp(inAtMin * inAtMax, 0.0, 1.0);
+
+    float4 color   = mix(uniforms.outOfRangeColor, uniforms.validRangeColor, t);
+    color.a        = uniforms.alpha;
+    return color;
+}
diff --git a/Sources/StructureKit/Metal/STKDepthBandOverlayRenderer.swift b/Sources/StructureKit/Metal/STKDepthBandOverlayRenderer.swift
@@ -0,0 +1,158 @@
+// StructureKit - A collection of extension utilities for Structure SDK
+// Copyright 2022 XRPro, LLC. All rights reserved.
+// http://structure.io
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+// * Neither the name of XRPro, LLC nor the names of its contributors may be
+//   used to endorse or promote products derived from this software without
+//   specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+
+import Metal
+import MetalKit
+import StructureKitCTypes
+
+public class STKDepthBandOverlayRenderer {
+  private var mtkView: MTKView
+  private var renderDepthState: MTLRenderPipelineState
+  private var textureDepth: MTLTexture?
+  private var vertexBuffer: MTLBuffer
+  private var samplerState: MTLSamplerState
+  private var device: MTLDevice
+  private var intr: STKIntrinsics?
+
+  private var validRangeMinMM: Float = 200.0
+  private var validRangeMaxMM: Float = 400.0
+  private var validRangeColor: simd_float4 = simd_float4(0,1,0,0.5) // Green
+  private var outOfRangeColor: simd_float4 = simd_float4(1,0,0,0.5) // Red
+  private var feather: Float = 40.0 // in mm
+
+  public init(view: MTKView, device: MTLDevice) {
+    mtkView = view
+    self.device = device
+
+    let library = STKMetalLibLoader.load(device: device)
+    let pipelineDepthDescriptor = MTLRenderPipelineDescriptor()
+    pipelineDepthDescriptor.sampleCount = 1
+    pipelineDepthDescriptor.colorAttachments[0].pixelFormat = mtkView.colorPixelFormat
+    pipelineDepthDescriptor.depthAttachmentPixelFormat = mtkView.depthStencilPixelFormat
+    pipelineDepthDescriptor.vertexFunction = library.makeFunction(name: "vertexDepthBandOverlay")
+    pipelineDepthDescriptor.fragmentFunction = library.makeFunction(name: "fragmentDepthBandOverlay")
+    // blending
+    pipelineDepthDescriptor.colorAttachments[0].isBlendingEnabled = true
+    pipelineDepthDescriptor.colorAttachments[0].rgbBlendOperation = .add
+    pipelineDepthDescriptor.colorAttachments[0].sourceRGBBlendFactor = .sourceAlpha
+    pipelineDepthDescriptor.colorAttachments[0].destinationRGBBlendFactor = .oneMinusSourceAlpha
+
+    renderDepthState = try! device.makeRenderPipelineState(descriptor: pipelineDepthDescriptor)
+
+    let vertexData = makeRectangularVertices()
+    let dataSize = vertexData.count * MemoryLayout.size(ofValue: vertexData[0])
+    vertexBuffer = device.makeBuffer(bytes: vertexData, length: dataSize, options: [])!
+
+    samplerState = makeDefaultSampler(device)
+  }
+
+  public func uploadColorTextureFromDepth(_ depthFrame: STKDepthFrame) {
+    if let texture = textureDepth {
+      if texture.width != depthFrame.width || texture.height != depthFrame.height {
+        textureDepth = nil  // invalidate texture
+      }
+    }
+
+    if textureDepth == nil {
+      let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
+        pixelFormat: .r32Float, width: Int(depthFrame.width), height: Int(depthFrame.height), mipmapped: false)
+      textureDescriptor.usage = MTLTextureUsage(
+        rawValue: MTLTextureUsage.renderTarget.rawValue | MTLTextureUsage.shaderRead.rawValue)
+      textureDepth = device.makeTexture(descriptor: textureDescriptor)
+    }
+
+    intr = depthFrame.intrinsics()
+    let depthMap = depthFrame.depthInMillimeters
+    assert(MemoryLayout<Float32>.size == MemoryLayout<Float>.size)
+
+    let bytesPerRow: Int = Int(depthFrame.width) * MemoryLayout<Float32>.stride
+    let region = MTLRegionMake2D(0, 0, Int(depthFrame.width), Int(depthFrame.height))
+    textureDepth?.replace(region: region, mipmapLevel: 0, withBytes: depthMap!, bytesPerRow: bytesPerRow)
+  }
+
+  public func renderDepthOverlay(
+    _ commandEncoder: MTLRenderCommandEncoder,
+    volumeSizeInMeters: simd_float3,
+    cameraPosition: float4x4,
+    textureOrientation: float4x4,
+    alpha: Float
+  ) {
+    guard let texture = textureDepth,
+      let intr = intr
+    else { return }
+
+    commandEncoder.pushDebugGroup("RenderDepthRangeOverlay")
+    commandEncoder.setRenderPipelineState(renderDepthState)
+
+    // rotate and mirror:
+    // move to [-0.5, 0.5] coordinates
+    // rotate
+    // mirror and apply scale to fix the ratio
+    // move back to [0, 1] coordinates
+    let projection =
+      float4x4.makeTranslation(0.5, 0.5, 0) * textureOrientation * float4x4.makeTranslation(-0.5, -0.5, 0)
+
+    let intrinsics = STKIntrinsicsMetal(
+      cx: intr.cx, cy: intr.cy, fx: intr.fx, fy: intr.fy, width: UInt32(texture.width), height: UInt32(texture.height))
+
+    var uniforms = STKUniformsDepthBandOverlay(
+      projection: projection,
+      cameraPose: cameraPosition,
+      cameraIntrinsics: intrinsics,
+      cubeModelInv: float4x4.makeScale(volumeSizeInMeters.x, volumeSizeInMeters.y, volumeSizeInMeters.z).inverse,
+      alpha: alpha,
+      validRangeMinMM: validRangeMinMM,
+      validRangeMaxMM: validRangeMaxMM,
+      validRangeColor: validRangeColor,
+      outOfRangeColor: outOfRangeColor,
+      feather: feather
+    )
+
+    commandEncoder.setVertexBuffer(vertexBuffer, offset: 0, index: 0)
+    commandEncoder.setVertexBytes(&uniforms, length: MemoryLayout<STKUniformsDepthBandOverlay>.stride, index: 1)
+
+    commandEncoder.setFragmentBytes(&uniforms, length: MemoryLayout<STKUniformsDepthBandOverlay>.stride, index: 0)
+    commandEncoder.setFragmentTexture(texture, index: 0)  // [[ texture(0) ]],
+    commandEncoder.setFragmentSamplerState(samplerState, index: 0)  // [[ sampler(0) ]]
+    commandEncoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 6, instanceCount: 1)
+    commandEncoder.popDebugGroup()
+  }
+  
+  public func configure(
+      validRangeMinMM: Float, validRangeMaxMM: Float,
+      validRangeColor: simd_float4 = simd_float4(0,1,0,0.5),
+      outOfRangeColor: simd_float4 = simd_float4(1,0,0,0.5),
+      feather: Float = 40.0
+  ) {
+      self.validRangeMinMM = validRangeMinMM
+      self.validRangeMaxMM = validRangeMaxMM
+      self.validRangeColor = validRangeColor
+      self.outOfRangeColor = outOfRangeColor
+      self.feather = feather
+  }
+}
diff --git a/Sources/StructureKit/Metal/STKDepthRenderer.swift b/Sources/StructureKit/Metal/STKDepthRenderer.swift
@@ -34,6 +34,7 @@ import StructureKitCTypes
 // 1. the depth frame
 // 2. the cube
 // 3. the depth overlay inside the cube
+
 public class STKDepthRenderer {
   private var mtkView: MTKView
   private var renderDepthState: MTLRenderPipelineState
@@ -52,7 +53,7 @@ public class STKDepthRenderer {
   private var _renderCubeState: MTLRenderPipelineState!
   private var _vertexCubeBuffer: MTLBuffer!
   private var _indexCubeBuffer: MTLBuffer!
-
+  
   public init(view: MTKView, device: MTLDevice) {
     mtkView = view
     self.device = device
@@ -298,5 +299,4 @@ public class STKDepthRenderer {
       indexBufferOffset: 0)
     commandEncoder.popDebugGroup()
   }
-
 }
diff --git a/Sources/StructureKit/Metal/STKMetalData.h b/Sources/StructureKit/Metal/STKMetalData.h
@@ -73,6 +73,20 @@ struct STKUniformsDepthOverlay
     float alpha;
 };
 
+struct STKUniformsDepthBandOverlay
+{
+    matrix_float4x4 projection;
+    matrix_float4x4 cameraPose;
+    struct STKIntrinsicsMetal cameraIntrinsics;
+    matrix_float4x4 cubeModelInv;
+    float alpha;
+    float validRangeMinMM;
+    float validRangeMaxMM;
+    vector_float4 validRangeColor;
+    vector_float4 outOfRangeColor;
+    float feather;
+};
+
 struct STKUniformsLine
 {
     matrix_float4x4 model;
diff --git a/Sources/StructureKit/Metal/STKMetalRenderer.swift b/Sources/StructureKit/Metal/STKMetalRenderer.swift
@@ -169,6 +169,7 @@ public class STKMetalRenderer: NSObject, STKRenderer {
   private var _arkitRenderer: STKARKitOverlayRenderer
   private var _anchorRenderer: STKLineRenderer
   private var _depthOverlayRenderer: STKDepthRenderer
+  private var _depthBandOverlayRenderer: STKDepthBandOverlayRenderer
   private var _meshRenderer: STKScanMeshRenderer
   private var _thickLineRenderer: STKMeshRendererThickLines
   
@@ -215,6 +216,7 @@ public class STKMetalRenderer: NSObject, STKRenderer {
     _arkitRenderer = STKARKitOverlayRenderer(view: view, device: device)
     _anchorRenderer = STKLineRenderer(view: view, device: device)
     _depthOverlayRenderer = STKDepthRenderer(view: view, device: device)
+    _depthBandOverlayRenderer = STKDepthBandOverlayRenderer(view: view, device: device)
     _meshRenderer = STKScanMeshRenderer(view: view, device: device)
     _thickLineRenderer = STKMeshRendererThickLines(view: view, device: device)
 
@@ -226,6 +228,7 @@ public class STKMetalRenderer: NSObject, STKRenderer {
 
   public func setDepthFrame(_ depthFrame: STKDepthFrame) {
     _depthOverlayRenderer.uploadColorTextureFromDepth(depthFrame)
+    _depthBandOverlayRenderer.uploadColorTextureFromDepth(depthFrame)
     depthCameraGLProjectionMatrix = float4x4(depthFrame.glProjectionMatrix())
   }
 
@@ -261,6 +264,16 @@ public class STKMetalRenderer: NSObject, STKRenderer {
   public func setDepthRenderingColors(_ baseColors: [simd_float4]) {
     _depthOverlayRenderer.depthRenderingColors = baseColors
   }
+  
+  public func configureDepthBandOverlay(
+    validRangeMinMM: Float = 0,
+    validRangeMaxMM: Float = 0,
+    validRangeColor: simd_float4 = simd_float4(0,1,0,0.5),
+    outOfRangeColor: simd_float4 = simd_float4(1,0,0,0.5),
+    feather: Float = 40.0)
+  {
+    _depthBandOverlayRenderer.configure(validRangeMinMM: validRangeMinMM, validRangeMaxMM: validRangeMaxMM, validRangeColor: validRangeColor, outOfRangeColor: outOfRangeColor, feather: feather)
+  }
 
   public func startRendering() {
     guard let commandBuffer = _commandQueue.makeCommandBuffer(),
@@ -332,6 +345,16 @@ public class STKMetalRenderer: NSObject, STKRenderer {
       textureOrientation: textureOrientation,
       alpha: alpha)
   }
+  
+  public func renderHighlightedDepthBand(cameraPose: simd_float4x4, alpha: Float, textureOrientation: simd_float4x4) {
+    guard let commandEncoder = _commandEncoder else { return }
+    _depthBandOverlayRenderer.renderDepthOverlay(
+      commandEncoder,
+      volumeSizeInMeters: _volumeSize,
+      cameraPosition: cameraPose,
+      textureOrientation: textureOrientation,
+      alpha: alpha)
+  }
 
   public func renderDepthFrame(orientation textureOrientation: simd_float4x4, range: simd_float2) {
     guard let commandEncoder = _commandEncoder else { return }
