feat: speed up point-object assignment with a bounding volume hierarchy

examples/src/utils/PointCloudObjectStylingUI.ts

@@ -23,10 +23,27 @@ export class PointCloudObjectStylingUI {
     const actions = {
       reset: () => {
         this._model.removeAllStyledObjectCollections();
+      },
+      randomColors: () => {
+        model.traverseStylableObjects((object) => {
+          const objectStyle: [number, number, number] = [
+            Math.floor(Math.random() * 255),
+            Math.floor(Math.random() * 255),
+            Math.floor(Math.random() * 255),
+          ];
+
+          const stylableObject = new AnnotationIdPointCloudObjectCollection([
+            object.annotationId,
+          ]);
+          model.assignStyledObjectCollection(stylableObject, {
+            color: objectStyle,
+          });
+        });
       }
     };
 
     uiFolder.add(actions, 'reset').name('Reset all styled objects');
+    uiFolder.add(actions, 'randomColors').name('Set random for objects');
   }
 
   private createObjectAppearanceUi(uiFolder: dat.GUI): () => PointCloudAppearance {

viewer/packages/pointclouds/src/bvh/BoundingVolumeHierarchy.test.ts

@@ -0,0 +1,74 @@
+/*!
+ * Copyright 2022 Cognite AS
+ */
+
+import { BoundingVolumeHierarchy } from './BoundingVolumeHierarchy';
+import { BvhElement } from './BvhElement';
+
+import * as SeededRandom from 'random-seed';
+
+import * as THREE from 'three';
+
+const random = SeededRandom.create('someseed');
+
+class BasicBvhElement implements BvhElement {
+  constructor(private readonly _box: THREE.Box3, private readonly _id: number) {}
+
+  getBox(): THREE.Box3 {
+    return this._box;
+  }
+
+  getId() {
+    return this._id;
+  }
+}
+
+function createId() {
+  return Math.floor(random.random() * 100000);
+}
+
+function createBoxes(): BasicBvhElement[] {
+  return new Array<number>(100).fill(0).map(_ => {
+    const min = new THREE.Vector3(random.random(), random.random(), random.random()).multiplyScalar(100);
+    const span = new THREE.Vector3(random.random(), random.random(), random.random()).multiplyScalar(50);
+    const box = new THREE.Box3(min, min.clone().add(span));
+
+    return new BasicBvhElement(box, createId());
+  });
+}
+
+function createPoints(): THREE.Vector3[] {
+  return new Array<number>(20).fill(0).map(_ => {
+    return new THREE.Vector3(random.random(), random.random(), random.random()).multiplyScalar(120);
+  });
+}
+
+describe(BoundingVolumeHierarchy.name, () => {
+  test('finds all intersecting elements for query points', () => {
+    const boxes = createBoxes();
+    const bvh = new BoundingVolumeHierarchy(boxes);
+    const points = createPoints();
+
+    let numBoxesFound = 0;
+
+    for (const point of points) {
+      const expectedBoxIds: number[] = [];
+
+      for (const box of boxes) {
+        if (box.getBox().containsPoint(point)) {
+          expectedBoxIds.push(box.getId());
+        }
+      }
+
+      const newBoxIds = bvh.findContainingElements(point).map(b => b.getId());
+
+      newBoxIds.sort();
+      expectedBoxIds.sort();
+
+      expect(newBoxIds).toEqual(expectedBoxIds);
+      numBoxesFound += newBoxIds.length;
+    }
+
+    expect(numBoxesFound).toBePositive();
+  });
+});

viewer/packages/pointclouds/src/bvh/BoundingVolumeHierarchy.ts

@@ -0,0 +1,28 @@
+/*!
+ * Copyright 2022 Cognite AS
+ */
+
+import { BvhElement } from './BvhElement';
+import { BvhNode } from './BvhNode';
+
+export class BoundingVolumeHierarchy<T extends BvhElement> {
+  private readonly _root: BvhNode<T>;
+
+  constructor(elements: T[]) {
+    this._root = new BvhNode(elements);
+  }
+
+  get root(): BvhNode<T> {
+    return this._root;
+  }
+
+  findContainingElements(point: THREE.Vector3): T[] {
+    const resultList = new Array<T>();
+    this._root.findContainingElements(point, resultList);
+    return resultList;
+  }
+
+  traverseContainingElements(point: THREE.Vector3, callback: (element: T) => void): void {
+    return this._root.traverseContainingElements(point, callback);
+  }
+}

viewer/packages/pointclouds/src/bvh/BvhElement.ts

@@ -0,0 +1,7 @@
+/*!
+ * Copyright 2022 Cognite AS
+ */
+
+export interface BvhElement {
+  getBox(): THREE.Box3;
+}

viewer/packages/pointclouds/src/bvh/BvhNode.ts

@@ -0,0 +1,87 @@
+/*!
+ * Copyright 2022 Cognite AS
+ */
+
+import { unionBoxes } from './unionBoxes';
+import * as THREE from 'three';
+
+import { BvhElement } from './BvhElement';
+
+import { findBestSplit } from './bvhUtils';
+
+const MAX_ELEMENTS_IN_LEAF = 1;
+
+export class BvhNode<T extends BvhElement> {
+  // Either _children or _elements is defined. Not both, not none.
+  private readonly _children: [BvhNode<T>, BvhNode<T>] | undefined;
+  private readonly _elements: T[] | undefined;
+
+  private readonly _boundingBox: THREE.Box3;
+
+  constructor(elements: T[]) {
+    if (elements.length <= MAX_ELEMENTS_IN_LEAF) {
+      this._elements = elements.slice();
+      this._boundingBox = unionBoxes(this._elements.map(e => e.getBox()));
+
+      return;
+    }
+
+    const [firstElements, lastElements] = findBestSplit(elements);
+
+    this._children = [new BvhNode<T>(firstElements), new BvhNode<T>(lastElements)];
+
+    this._boundingBox = this._children[0].boundingBox.clone().union(this._children[1].boundingBox);
+  }
+
+  get boundingBox(): THREE.Box3 {
+    return this._boundingBox;
+  }
+
+  get children(): [BvhNode<T>, BvhNode<T>] | undefined {
+    return this._children;
+  }
+
+  get elements(): T[] | undefined {
+    return this._elements;
+  }
+
+  findContainingElements(point: THREE.Vector3, resultList: T[]): void {
+    if (this._elements) {
+      for (const element of this._elements) {
+        if (element.getBox().containsPoint(point)) {
+          resultList.push(element);
+        }
+      }
+
+      return;
+    }
+
+    if (this._children![0].boundingBox.containsPoint(point)) {
+      this._children![0].findContainingElements(point, resultList);
+    }
+
+    if (this._children![1].boundingBox.containsPoint(point)) {
+      this._children![1].findContainingElements(point, resultList);
+    }
+  }
+
+  traverseContainingElements(point: THREE.Vector3, callback: (element: T) => void): void {
+    if (this._elements) {
+      for (const element of this._elements) {
+        if (element.getBox().containsPoint(point)) {
+          callback(element);
+        }
+      }
+
+      return;
+    }
+
+    if (this._children![0].boundingBox.containsPoint(point)) {
+      this._children![0].traverseContainingElements(point, callback);
+    }
+
+    if (this._children![1].boundingBox.containsPoint(point)) {
+      this._children![1].traverseContainingElements(point, callback);
+    }
+  }
+}

viewer/packages/pointclouds/src/bvh/bvhUtils.ts

@@ -0,0 +1,86 @@
+/*!
+ * Copyright 2022 Cognite AS
+ */
+
+import * as THREE from 'three';
+import { BvhElement } from './BvhElement';
+
+import { unionBoxes } from './unionBoxes';
+
+const computeVectorSeparatismVars = { vec: new THREE.Vector3() };
+
+function computeVectorSeparatism(min: THREE.Vector3, max: THREE.Vector3): number {
+  const { vec: diff } = computeVectorSeparatismVars;
+
+  diff.copy(min).sub(max);
+
+  return Math.max(diff.x, Math.max(diff.y, diff.z));
+}
+
+/**
+ * Computes "separatism" - the degree to which the boxes are separated (will be negative if boxes are intersecting)
+ */
+function computeBoxSeparatism(a: THREE.Box3, b: THREE.Box3): number {
+  return Math.max(computeVectorSeparatism(a.min, b.max), computeVectorSeparatism(b.min, a.max));
+}
+
+/**
+ * Simple box sorting function that just sorts the boxes by their value in their
+ * `axisIndex` component of the `min` vector
+ */
+function elementAxisCompare(e0: BvhElement, e1: BvhElement, axisIndex: number): number {
+  return e0.getBox().min.getComponent(axisIndex) - e1.getBox().min.getComponent(axisIndex);
+}
+
+function copySortedByAxis<T extends BvhElement>(elements: T[], axisIndex: number): T[] {
+  const elementsCopy = elements.slice();
+  elementsCopy.sort((e0, e1) => elementAxisCompare(e0, e1, axisIndex));
+  return elementsCopy;
+}
+
+function splitByMiddle<T>(elements: T[]): [T[], T[]] {
+  const midIndex = Math.floor(elements.length / 2);
+  return [elements.slice(0, midIndex), elements.slice(midIndex, elements.length)];
+}
+
+/**
+ * Returns how well sorting the `elements` along the specified axis "separates" the
+ * bounding boxes of the two halves of the list. A good separation heuristically means
+ * it's a good sorting for splitting the elements at a BVH node
+ */
+function evaluateSeparatingAxis(elements: BvhElement[], axisIndex: number): number {
+  if (elements.length < 2) {
+    throw Error('Too few elements provided to separation evaluation function, need at least two');
+  }
+
+  const sortedElements = copySortedByAxis(elements, axisIndex);
+  const boxes = sortedElements.map(e => e.getBox());
+
+  const [firstBoxes, lastBoxes] = splitByMiddle(boxes);
+
+  const box0 = unionBoxes(firstBoxes);
+  const box1 = unionBoxes(lastBoxes);
+
+  return computeBoxSeparatism(box0, box1);
+}
+
+/**
+ * Splits the list into two (somewhat) equally sized parts which is deemed to best
+ * separates the boxes associated with the elements
+ */
+export function findBestSplit<T extends BvhElement>(elements: T[]): [T[], T[]] {
+  let bestSeparatism = -Infinity,
+    bestAxisIndex = 0;
+
+  for (let i = 0; i < 3; i++) {
+    const separatismValue = evaluateSeparatingAxis(elements, i);
+
+    if (separatismValue > bestSeparatism) {
+      bestAxisIndex = i;
+      bestSeparatism = separatismValue;
+    }
+  }
+
+  const sortedElements = copySortedByAxis(elements, bestAxisIndex);
+  return splitByMiddle(sortedElements);
+}

viewer/packages/pointclouds/src/bvh/unionBoxes.ts

@@ -0,0 +1,17 @@
+/*!
+ * Copyright 2022 Cognite AS
+ */
+
+import * as THREE from 'three';
+
+export function unionBoxes(boxes: THREE.Box3[], out?: THREE.Box3): THREE.Box3 {
+  out = out ?? new THREE.Box3();
+  out.makeEmpty();
+
+  boxes.forEach(inputBox => {
+    out!.expandByPoint(inputBox.max);
+    out!.expandByPoint(inputBox.min);
+  });
+
+  return out;
+}

viewer/packages/pointclouds/src/factory/CdfPointCloudFactory.test.ts

@@ -11,6 +11,8 @@ import { PointCloudMetadata } from '../PointCloudMetadata';
 import { Potree, PointCloudOctree, PointCloudMaterial } from '../potree-three-loader';
 import { PotreeNodeWrapper } from '../PotreeNodeWrapper';
 
+import * as THREE from 'three';
+
 const dummyAnnotationsResponse = {
   items: [
     {
@@ -75,18 +77,15 @@ describe(CdfPointCloudFactory.name, () => {
   });
 
   test('contains right geometries for annotations provided by SDK', async () => {
-    const expectedContainedPoints: [number, number, number][] = [
-      [-0.03, 0.1, -500],
-      [0.4, -0.4, 0]
-    ];
-    const expectedUncontainedPoints: [number, number, number][] = [
-      [1, 1, 1],
-      [300, 300, 300]
+    const expectedContainedPoints: THREE.Vector3[] = [
+      new THREE.Vector3(-0.03, 0.1, -500),
+      new THREE.Vector3(0.4, -0.4, 0)
     ];
+    const expectedUncontainedPoints: THREE.Vector3[] = [new THREE.Vector3(1, 1, 1), new THREE.Vector3(300, 300, 300)];
 
     const shapes = model.stylableObjects.map(obj => obj.stylableObject.shape);
 
-    function containedInAnyShape(p: [number, number, number]): boolean {
+    function containedInAnyShape(p: THREE.Vector3): boolean {
       let contained = false;
       for (const shape of shapes) {
         contained ||= shape.containsPoint(p);

viewer/packages/pointclouds/src/factory/CdfPointCloudFactory.ts

@@ -16,11 +16,12 @@ import { DEFAULT_POINT_CLOUD_METADATA_FILE } from '../constants';
 import { CogniteClient } from '@cognite/sdk';
 
 import { Box } from '../styling/shapes/Box';
-import { createInvertedRevealTransformationFromCdfTransformation } from '../styling/shapes/linalg';
 import { Cylinder } from '../styling/shapes/Cylinder';
 import { IShape } from '../styling/shapes/IShape';
 import { PointCloudFactory } from '../IPointCloudFactory';
 
+import * as THREE from 'three';
+
 export class CdfPointCloudFactory implements PointCloudFactory {
   private readonly _potreeInstance: Potree;
   private readonly _sdkClient: CogniteClient;
@@ -36,11 +37,15 @@ export class CdfPointCloudFactory implements PointCloudFactory {
 
   private annotationGeometryToLocalGeometry(geometry: any): IShape {
     if (geometry.box) {
-      return new Box(createInvertedRevealTransformationFromCdfTransformation({ data: geometry.box.matrix }));
+      return new Box(new THREE.Matrix4().fromArray(geometry.box.matrix).transpose().invert());
     }
 
     if (geometry.cylinder) {
-      return new Cylinder(geometry.cylinder.centerA, geometry.cylinder.centerB, geometry.cylinder.radius);
+      return new Cylinder(
+        new THREE.Vector3().fromArray(geometry.cylinder.centerA),
+        new THREE.Vector3().fromArray(geometry.cylinder.centerB),
+        geometry.cylinder.radius
+      );
     }
 
     throw Error('Annotation geometry type not recognized');