Allow dynamic change of natively rendered layer (by re-applying transforms)

frontend/javascripts/libs/estimate_affine.ts

@@ -1,4 +1,5 @@
 import _ from "lodash";
+import { Matrix4x4 } from "mjs";
 import { Matrix, solve } from "ml-matrix";
 import { Vector3 } from "oxalis/constants";
 
@@ -28,11 +29,13 @@ export default function estimateAffine(sourcePoints: Vector3[], targetPoints: Ve
   const xMatrix = solve(A, b);
   const x = xMatrix.to1DArray();
   const error = Matrix.sub(b, new Matrix(A).mmul(xMatrix)).to1DArray();
-  console.log(
-    "Affine estimation error: ",
-    error,
-    `(mean=${_.mean(error.map((el) => Math.abs(el)))})`,
-  );
+  if (!process.env.IS_TESTING) {
+    console.log(
+      "Affine estimation error: ",
+      error,
+      `(mean=${_.mean(error.map((el) => Math.abs(el)))})`,
+    );
+  }
 
   const affineMatrix = new Matrix([
     [x[0], x[1], x[2], x[3]],
@@ -43,3 +46,7 @@ export default function estimateAffine(sourcePoints: Vector3[], targetPoints: Ve
 
   return new Matrix(affineMatrix);
 }
+
+export function estimateAffineMatrix4x4(sourcePoints: Vector3[], targetPoints: Vector3[]) {
+  return estimateAffine(sourcePoints, targetPoints).to1DArray() as any as Matrix4x4;
+}

frontend/javascripts/libs/mjs.ts

@@ -75,6 +75,9 @@ const M4x4 = {
   // and also returns Array<Vector3>
   transformVectorsAffine(m: Matrix4x4, _points: Vector3[]): Vector3[] {
     const points: Array<Array<number>> = _points as any as Array<Array<number>>;
+    if (!Array.isArray(_points[0])) {
+      throw new Error("transformVectorsAffine doesn't support typed arrays at the moment.");
+    }
     // @ts-ignore
     return chunk3(M4x4.transformPointsAffine(m, _.flatten(points)));
   },

frontend/javascripts/libs/multi_key_map.ts

@@ -0,0 +1,35 @@
+type RecursiveMap<K, V> = Map<K, V | RecursiveMap<K, V>>;
+
+export default class MultiKeyMap<E, V, K extends E[]> {
+  map: RecursiveMap<E, V> = new Map<E, RecursiveMap<E, V>>();
+  set(keys: K, value: V) {
+    let currentMap = this.map;
+    for (const [index, key] of keys.entries()) {
+      if (index < keys.length - 1) {
+        let foundMap = currentMap.get(key) as RecursiveMap<E, V> | undefined;
+        if (foundMap == null) {
+          foundMap = new Map<E, RecursiveMap<E, V>>();
+          currentMap.set(key, foundMap);
+        }
+        currentMap = foundMap;
+      } else {
+        currentMap.set(key, value);
+      }
+    }
+  }
+  get(keys: K): V | undefined {
+    let currentMap = this.map;
+    for (const [index, key] of keys.entries()) {
+      if (index < keys.length - 1) {
+        const foundMap = currentMap.get(key);
+        if (foundMap === undefined) {
+          return undefined;
+        }
+        currentMap = foundMap as RecursiveMap<E, V>;
+      } else {
+        return currentMap.get(key) as V | undefined;
+      }
+    }
+    throw new Error("Unexpected code reached.");
+  }
+}

frontend/javascripts/libs/thin_plate_spline.ts

@@ -1,6 +1,7 @@
 import _ from "lodash";
 import { Matrix, solve } from "ml-matrix";
 import { Vector3 } from "oxalis/constants";
+import { V3 } from "./mjs";
 
 class TPS1d {
   // This class accepts 3-dimensional control points
@@ -110,13 +111,24 @@ export default class TPS3D {
   tpsY = new TPS1d();
   tpsZ = new TPS1d();
 
-  constructor(sourcePoints: Vector3[], targetPoints: Vector3[]) {
+  unscaledSourcePoints: Vector3[];
+  unscaledTargetPoints: Vector3[];
+  scale: Vector3;
+
+  constructor(unscaledSourcePoints: Vector3[], unscaledTargetPoints: Vector3[], scale: Vector3) {
+    const sourcePoints = unscaledSourcePoints.map((point) => V3.scale3(point, scale, [0, 0, 0]));
+    const targetPoints = unscaledTargetPoints.map((point) => V3.scale3(point, scale, [0, 0, 0]));
+
     const [cps, offsetX, offsetY, offsetZ] = this.getControlPointsWithOffsets(
       sourcePoints,
       targetPoints,
       true,
     );
 
+    this.unscaledSourcePoints = sourcePoints;
+    this.unscaledTargetPoints = targetPoints;
+    this.scale = scale;
+
     this.tpsX.fit(offsetX, cps);
     this.tpsY.fit(offsetY, cps);
     this.tpsZ.fit(offsetZ, cps);

frontend/javascripts/libs/utils.ts

@@ -1149,3 +1149,19 @@ export function minValue(array: Array<number>): number {
   }
   return value;
 }
+
+/*
+ * Iterates over arbitrary objects recursively and calls the callback function.
+ */
+export const deepIterate = (obj: any, callback: (val: unknown) => void) => {
+  if (obj == null) {
+    return;
+  }
+  Object.keys(obj).forEach((key: any) => {
+    callback(obj[key]);
+
+    if (typeof obj[key] === "object" && obj[key] !== null) {
+      deepIterate(obj[key], callback);
+    }
+  });
+};

frontend/javascripts/oxalis/api/api_latest.ts

@@ -78,6 +78,7 @@ import {
   getResolutionInfo,
   getVisibleSegmentationLayer,
   getMappingInfo,
+  flatToNestedMatrix,
 } from "oxalis/model/accessors/dataset_accessor";
 import {
   getPosition,
@@ -2038,7 +2039,7 @@ class DataApi {
    *
    * @example
    *
-   * api.data._setLayerTransforms(
+   * api.data._setAffineLayerTransforms(
    *   "C555_DIAMOND_2f",
    *   new Float32Array([
    *     0.03901274364025348, -0.08498337289603758, 0.00782446404039791, 555.7948181512004,
@@ -2048,8 +2049,15 @@ class DataApi {
    *   ]),
    * );
    */
-  _setLayerTransforms(layerName: string, transforms: Matrix4x4) {
-    Store.dispatch(setLayerTransformsAction(layerName, Array.from(transforms) as Vector16));
+  _setAffineLayerTransforms(layerName: string, transforms: Matrix4x4) {
+    const coordinateTransforms = [
+      {
+        type: "affine" as const,
+        matrix: flatToNestedMatrix(Array.from(transforms) as Vector16),
+      },
+    ];
+
+    Store.dispatch(setLayerTransformsAction(layerName, coordinateTransforms));
   }
 
   /*

frontend/javascripts/oxalis/controller/scene_controller.ts

@@ -458,7 +458,11 @@ class SceneController {
           isHighlighted: false,
         });
         bbCube.getMeshes().forEach((mesh) => {
-          const transformMatrix = getTransformsForLayerOrNull(dataset, layer)?.affineMatrix;
+          const transformMatrix = getTransformsForLayerOrNull(
+            dataset,
+            layer,
+            state.datasetConfiguration.nativelyRenderedLayerName,
+          )?.affineMatrix;
           if (transformMatrix) {
             const matrix = new THREE.Matrix4();
             // @ts-ignore

frontend/javascripts/oxalis/geometries/materials/plane_material_factory.ts

@@ -31,6 +31,7 @@ import {
   getTransformsForLayer,
   getResolutionInfoByLayer,
   getResolutionInfo,
+  getTransformsPerLayer,
 } from "oxalis/model/accessors/dataset_accessor";
 import {
   getActiveMagIndicesForLayers,
@@ -226,6 +227,9 @@ class PlaneMaterialFactory {
     this.uniforms.activeMagIndices = {
       value: Object.values(activeMagIndices),
     };
+    const nativelyRenderedLayerName =
+      Store.getState().datasetConfiguration.nativelyRenderedLayerName;
+    const dataset = Store.getState().dataset;
     for (const dataLayer of Model.getAllLayers()) {
       const layerName = sanitizeName(dataLayer.name);
 
@@ -241,14 +245,18 @@ class PlaneMaterialFactory {
       this.uniforms[`${layerName}_unrenderable`] = {
         value: 0,
       };
-      const dataset = Store.getState().dataset;
       const layer = getLayerByName(dataset, dataLayer.name);
 
       this.uniforms[`${layerName}_transform`] = {
-        value: invertAndTranspose(getTransformsForLayer(dataset, layer).affineMatrix),
+        value: invertAndTranspose(
+          getTransformsForLayer(dataset, layer, nativelyRenderedLayerName).affineMatrix,
+        ),
       };
       this.uniforms[`${layerName}_has_transform`] = {
-        value: !_.isEqual(getTransformsForLayer(dataset, layer).affineMatrix, Identity4x4),
+        value: !_.isEqual(
+          getTransformsForLayer(dataset, layer, nativelyRenderedLayerName).affineMatrix,
+          Identity4x4,
+        ),
       };
     }
 
@@ -431,8 +439,9 @@ class PlaneMaterialFactory {
           // isn't relevant which is why it can default to [1, 1, 1].
 
           let representativeMagForVertexAlignment: Vector3 = [Infinity, Infinity, Infinity];
+          const state = Store.getState();
           for (const [layerName, activeMagIndex] of Object.entries(activeMagIndices)) {
-            const layer = getLayerByName(Store.getState().dataset, layerName);
+            const layer = getLayerByName(state.dataset, layerName);
             const resolutionInfo = getResolutionInfo(layer.resolutions);
             // If the active mag doesn't exist, a fallback mag is likely rendered. Use that
             // to determine a representative mag.
@@ -443,7 +452,11 @@ class PlaneMaterialFactory {
                 : null;
 
             const hasTransform = !_.isEqual(
-              getTransformsForLayer(Store.getState().dataset, layer).affineMatrix,
+              getTransformsForLayer(
+                state.dataset,
+                layer,
+                state.datasetConfiguration.nativelyRenderedLayerName,
+              ).affineMatrix,
               Identity4x4,
             );
             if (!hasTransform && suitableMag) {
@@ -780,27 +793,37 @@ class PlaneMaterialFactory {
 
     this.storePropertyUnsubscribers.push(
       listenToStoreProperty(
-        (storeState) => storeState.dataset.dataSource.dataLayers,
-        (layers) => {
+        (storeState) =>
+          getTransformsPerLayer(
+            storeState.dataset,
+            storeState.datasetConfiguration.nativelyRenderedLayerName,
+          ),
+        (transformsPerLayer) => {
           this.scaledTpsInvPerLayer = {};
+          const state = Store.getState();
+          const layers = state.dataset.dataSource.dataLayers;
           for (let layerIdx = 0; layerIdx < layers.length; layerIdx++) {
             const layer = layers[layerIdx];
             const name = sanitizeName(layer.name);
-            const transforms = getTransformsForLayer(Store.getState().dataset, layer);
+            const transforms = transformsPerLayer[layer.name];
             const { affineMatrix } = transforms;
             const scaledTpsInv =
               transforms.type === "thin_plate_spline" ? transforms.scaledTpsInv : null;
 
             if (scaledTpsInv) {
               this.scaledTpsInvPerLayer[name] = scaledTpsInv;
+            } else {
+              delete this.scaledTpsInvPerLayer[name];
             }
 
             this.uniforms[`${name}_transform`].value = invertAndTranspose(affineMatrix);
             const hasTransform = !_.isEqual(affineMatrix, Identity4x4);
+            console.log(`${name}_has_transform`, hasTransform);
             this.uniforms[`${name}_has_transform`] = {
               value: hasTransform,
             };
           }
+          this.recomputeShaders();
         },
         true,
       ),
@@ -855,6 +878,9 @@ class PlaneMaterialFactory {
   }
 
   recomputeShaders = _.throttle(() => {
+    if (this.material == null) {
+      return;
+    }
     const [newFragmentShaderCode, additionalUniforms] = this.getFragmentShaderWithUniforms();
     for (const [name, value] of Object.entries(additionalUniforms)) {
       this.uniforms[name] = value;