Speed up mapping creation/activation (#4103)

* speed up mapping creation/activation

* add api function

* add screenshot test for mappings

* update changelog

CHANGELOG.md

@@ -16,6 +16,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.md).
 - When holding CTRL while toggling the visibility of a layer, that layer will be made exclusively visible. This change makes it easier to quickly compare different data layers against each other. [#4061](https://github.com/scalableminds/webknossos/pull/4061)
 
 ### Changed
+- Heavily improved mapping creation/activation performance. [#4103](https://github.com/scalableminds/webknossos/pull/4103)
 - The NML parser now rounds floating point values in node coordinates. [#4045](https://github.com/scalableminds/webknossos/pull/4045)
 
 ### Fixed

frontend/javascripts/oxalis/api/api_latest.js

@@ -704,6 +704,9 @@ class DataApi {
       setMappingAction(
         "<custom mapping>",
         _.clone(mapping),
+        // Object.keys is sorted for numerical keys according to the spec:
+        // http://www.ecma-international.org/ecma-262/6.0/#sec-ordinary-object-internal-methods-and-internal-slots-ownpropertykeys
+        Object.keys(mapping).map(x => parseInt(x, 10)),
         options.colors,
         options.hideUnmappedIds,
       ),
@@ -736,6 +739,14 @@ class DataApi {
     return this.model.getSegmentationLayer().setActiveMapping(mappingName);
   }
 
+  /**
+   * Returns whether a mapping is currently enabled.
+   *
+   */
+  isMappingEnabled(): boolean {
+    return Store.getState().temporaryConfiguration.activeMapping.isMappingEnabled;
+  }
+
   /**
    * Returns the bounding box for a given layer name.
    */

frontend/javascripts/oxalis/api/api_v2.js

@@ -553,7 +553,13 @@ class DataApi {
     if (layerName !== segmentationLayerName) {
       throw new Error(messages["mapping.unsupported_layer"]);
     }
-    Store.dispatch(setMappingAction("<custom mapping>", _.clone(mapping)));
+    Store.dispatch(
+      setMappingAction(
+        "<custom mapping>",
+        _.clone(mapping),
+        Object.keys(mapping).map(x => parseInt(x, 10)),
+      ),
+    );
   }
 
   /**

frontend/javascripts/oxalis/default_state.js

@@ -81,6 +81,7 @@ const defaultState: OxalisState = {
     activeMapping: {
       mappingName: null,
       mapping: null,
+      mappingKeys: null,
       mappingColors: null,
       hideUnmappedIds: false,
       isMappingEnabled: false,

frontend/javascripts/oxalis/model/actions/settings_actions.js

@@ -51,6 +51,7 @@ type SetMappingAction = {
   type: "SET_MAPPING",
   mappingName: ?string,
   mapping: ?Mapping,
+  mappingKeys: ?Array<number>,
   mappingColors: ?Array<number>,
   hideUnmappedIds: ?boolean,
 };
@@ -145,12 +146,14 @@ export const setMappingEnabledAction = (isMappingEnabled: boolean): SetMappingEn
 export const setMappingAction = (
   mappingName: ?string,
   mapping: ?Mapping,
+  mappingKeys: ?Array<number>,
   mappingColors: ?Array<number>,
   hideUnmappedIds: ?boolean,
 ): SetMappingAction => ({
   type: "SET_MAPPING",
   mappingName,
   mapping,
+  mappingKeys,
   mappingColors,
   hideUnmappedIds,
 });

frontend/javascripts/oxalis/model/bucket_data_handling/mappings.js

@@ -94,9 +94,15 @@ class Mappings {
       // class will get the first color, and so on
       const assignNewIds = mappingColors != null && mappingColors.length > 0;
       await progressCallback(false, "Building mapping structure...");
-      const mappingObject = this.buildMappingObject(mappingName, fetchedMappings, assignNewIds);
+      const [mappingObject, mappingKeys] = this.buildMappingObject(
+        mappingName,
+        fetchedMappings,
+        assignNewIds,
+      );
       await progressCallback(false, "Applying mapping...");
-      Store.dispatch(setMappingAction(mappingName, mappingObject, mappingColors, hideUnmappedIds));
+      Store.dispatch(
+        setMappingAction(mappingName, mappingObject, mappingKeys, mappingColors, hideUnmappedIds),
+      );
     }
   }
 
@@ -136,8 +142,11 @@ class Mappings {
     mappingName: string,
     fetchedMappings: APIMappings,
     assignNewIds: boolean,
-  ): Mapping {
+  ): [Mapping, Array<number>] {
     const mappingObject: Mapping = {};
+    // Performance optimization: Object.keys(...) is slow for large objects
+    // keeping track of the keys in a separate array is ~5x faster
+    const mappingKeys = [];
 
     const maxId = this.getLargestSegmentId() + 1;
     // Initialize to the next multiple of 256 that is larger than maxId
@@ -152,12 +161,14 @@ class Mappings {
           const mappedId = mappingObject[minId] || minId;
           for (const id of mappingClass) {
             mappingObject[id] = mappedId;
+            mappingKeys.push(id);
           }
           newMappedId++;
         }
       }
     }
-    return mappingObject;
+    mappingKeys.sort((a, b) => a - b);
+    return [mappingObject, mappingKeys];
   }
 
   getMappingChain(mappingName: string, fetchedMappings: APIMappings): Array<string> {
@@ -198,7 +209,8 @@ class Mappings {
     listenToStoreProperty(
       state => state.temporaryConfiguration.activeMapping.mapping,
       mapping => {
-        this.updateMappingTextures(mapping);
+        const { mappingKeys } = Store.getState().temporaryConfiguration.activeMapping;
+        this.updateMappingTextures(mapping, mappingKeys);
       },
     );
 
@@ -224,50 +236,49 @@ class Mappings {
     );
   }
 
-  async updateMappingTextures(mapping: ?Mapping): Promise<void> {
-    if (mapping == null) return;
+  async updateMappingTextures(mapping: ?Mapping, mappingKeys: ?Array<number>): Promise<void> {
+    if (mapping == null || mappingKeys == null) return;
     const progressCallback = this.progressCallback;
 
     await progressCallback(false, "Create mapping texture...");
     console.time("Time to create mapping texture");
-    // $FlowFixMe Flow chooses the wrong library definition, because it doesn't seem to know that Object.keys always returns strings and throws an error
-    const keys = Uint32Array.from(Object.keys(mapping), x => parseInt(x, 10));
-    keys.sort((a, b) => a - b);
+
+    const mappingSize = mappingKeys.length;
+    // The typed arrays need to be padded with 0s so that their length is a multiple of MAPPING_TEXTURE_WIDTH
+    const paddedLength =
+      mappingSize + MAPPING_TEXTURE_WIDTH - (mappingSize % MAPPING_TEXTURE_WIDTH);
+
+    const keys = new Uint32Array(paddedLength);
+    const values = new Uint32Array(paddedLength);
+
+    keys.set(mappingKeys);
     // $FlowFixMe Flow doesn't recognize that mapping cannot be null or undefined :/
-    const values = Uint32Array.from(keys, key => mapping[key]);
+    values.set(mappingKeys.map(key => mapping[key]));
+
     // Instantiate the Uint8Arrays with the array buffer from the Uint32Arrays, so that each 32-bit value is converted
     // to four 8-bit values correctly
     const uint8Keys = new Uint8Array(keys.buffer);
     const uint8Values = new Uint8Array(values.buffer);
-    // The typed arrays need to be padded with 0s so that their length is a multiple of MAPPING_TEXTURE_WIDTH
-    const paddedLength =
-      keys.length + MAPPING_TEXTURE_WIDTH - (keys.length % MAPPING_TEXTURE_WIDTH);
-    // The length of typed arrays cannot be changed, so we need to create new ones with the correct length
-    const uint8KeysPadded = new Uint8Array(paddedLength * 4);
-    uint8KeysPadded.set(uint8Keys);
-    const uint8ValuesPadded = new Uint8Array(paddedLength * 4);
-    uint8ValuesPadded.set(uint8Values);
     console.timeEnd("Time to create mapping texture");
 
-    const mappingSize = keys.length;
     if (mappingSize > MAPPING_TEXTURE_WIDTH ** 2) {
       throw new Error(messages["mapping.too_big"]);
     }
 
     await progressCallback(false, "Copy mapping data to GPU...");
     this.mappingLookupTexture.update(
-      uint8KeysPadded,
+      uint8Keys,
       0,
       0,
       MAPPING_TEXTURE_WIDTH,
-      uint8KeysPadded.length / MAPPING_TEXTURE_WIDTH / 4,
+      uint8Keys.length / MAPPING_TEXTURE_WIDTH / 4,
     );
     this.mappingTexture.update(
-      uint8ValuesPadded,
+      uint8Values,
       0,
       0,
       MAPPING_TEXTURE_WIDTH,
-      uint8ValuesPadded.length / MAPPING_TEXTURE_WIDTH / 4,
+      uint8Values.length / MAPPING_TEXTURE_WIDTH / 4,
     );
     await progressCallback(true, "Mapping successfully applied.");
     Store.dispatch(setMappingEnabledAction(true));

frontend/javascripts/oxalis/model/reducers/settings_reducer.js

@@ -153,8 +153,9 @@ function SettingsReducer(state: OxalisState, action: Action): OxalisState {
     case "SET_MAPPING": {
       return updateActiveMapping(state, {
         mappingName: action.mappingName,
-        mappingSize: action.mapping != null ? _.size(action.mapping) : 0,
+        mappingSize: action.mappingKeys != null ? action.mappingKeys.length : 0,
         mapping: action.mapping,
+        mappingKeys: action.mappingKeys,
         mappingColors: action.mappingColors,
         hideUnmappedIds: action.hideUnmappedIds || false,
       });

frontend/javascripts/oxalis/store.js

@@ -280,6 +280,7 @@ export type TemporaryConfiguration = {
   +activeMapping: {
     +mappingName: ?string,
     +mapping: ?Mapping,
+    +mappingKeys: ?Array<number>,
     +mappingColors: ?Array<number>,
     +hideUnmappedIds: boolean,
     +isMappingEnabled: boolean,

frontend/javascripts/oxalis/view/settings/dataset_settings_view.js

@@ -19,7 +19,7 @@ import {
 import { findDataPositionForLayer } from "admin/admin_rest_api";
 import { getGpuFactorsWithLabels } from "oxalis/model/bucket_data_handling/data_rendering_logic";
 import { getMaxZoomValueForResolution } from "oxalis/model/accessors/flycam_accessor";
-import { hasSegmentation, isRgb } from "oxalis/model/accessors/dataset_accessor";
+import { hasSegmentation, getElementClass } from "oxalis/model/accessors/dataset_accessor";
 import { setPositionAction, setZoomStepAction } from "oxalis/model/actions/flycam_actions";
 import {
   updateDatasetSettingAction,
@@ -99,7 +99,7 @@ class DatasetSettings extends React.PureComponent<DatasetSettingsProps> {
     layerIndex: number,
     isLastLayer: boolean,
   ) => {
-    const isRGB = isRgb(this.props.dataset, layerName);
+    const elementClass = getElementClass(this.props.dataset, layerName);
     const { brightness, contrast, alpha, color, isDisabled } = layer;
     return (
       <div key={layerName}>
@@ -130,9 +130,7 @@ class DatasetSettings extends React.PureComponent<DatasetSettingsProps> {
               </div>
             </Tooltip>
             <span style={{ fontWeight: 700 }}>{layerName}</span>
-            <Tag style={{ cursor: "default", marginLeft: 8 }} color={isRGB && "#1890ff"}>
-              {isRGB ? "24-bit" : "8-bit"} Layer
-            </Tag>
+            <Tag style={{ cursor: "default", marginLeft: 8 }}>{elementClass} Layer</Tag>
             {this.getFindDataButton(layerName, isDisabled)}
           </Col>
         </Row>

frontend/javascripts/test/puppeteer/dataset_rendering.screenshot.js

@@ -7,7 +7,11 @@ import path from "path";
 import puppeteer, { type Browser } from "puppeteer";
 
 import { compareScreenshot } from "./screenshot_helpers";
-import { screenshotDataset, DEV_AUTH_TOKEN } from "./dataset_rendering_helpers";
+import {
+  screenshotDataset,
+  screenshotDatasetWithMapping,
+  DEV_AUTH_TOKEN,
+} from "./dataset_rendering_helpers";
 
 process.on("unhandledRejection", (err, promise) => {
   console.error("Unhandled rejection (promise: ", promise, ", reason: ", err, ").");
@@ -115,6 +119,13 @@ async function withRetry(
   }
 }
 
+function isPixelEquivalent(changedPixels, width, height) {
+  // There may be a difference of 0.1 %
+  const allowedThreshold = 0.1 / 100;
+  const allowedChangedPixel = allowedThreshold * width * height;
+  return changedPixels < allowedChangedPixel;
+}
+
 datasetNames.map(async datasetName => {
   test.serial(`it should render dataset ${datasetName} correctly`, async t => {
     await withRetry(
@@ -136,10 +147,7 @@ datasetNames.map(async datasetName => {
           datasetName,
         );
 
-        // There may be a difference of 0.1 %
-        const allowedThreshold = 0.1 / 100;
-        const allowedChangedPixel = allowedThreshold * width * height;
-        return changedPixels < allowedChangedPixel;
+        return isPixelEquivalent(changedPixels, width, height);
       },
       condition => {
         t.true(
@@ -151,6 +159,38 @@ datasetNames.map(async datasetName => {
   });
 });
 
+test.serial("it should render a dataset with mappings correctly", async t => {
+  const datasetName = "ROI2017_wkw";
+  const mappingName = "axons";
+  await withRetry(
+    3,
+    async () => {
+      const datasetId = { name: datasetName, owningOrganization: "Connectomics_Department" };
+      const { screenshot, width, height } = await screenshotDatasetWithMapping(
+        await getNewPage(t.context.browser),
+        URL,
+        datasetId,
+        mappingName,
+      );
+      const changedPixels = await compareScreenshot(
+        screenshot,
+        width,
+        height,
+        BASE_PATH,
+        `${datasetName}_with_mapping_${mappingName}`,
+      );
+
+      return isPixelEquivalent(changedPixels, width, height);
+    },
+    condition => {
+      t.true(
+        condition,
+        `Dataset with name: "${datasetName}" and mapping: "${mappingName}" does not look the same.`,
+      );
+    },
+  );
+});
+
 test.afterEach(async t => {
   await t.context.browser.close();
 });

frontend/javascripts/test/puppeteer/dataset_rendering_helpers.js

@@ -40,14 +40,41 @@ export async function screenshotDataset(
   if (optionalDatasetConfigOverride != null) {
     await updateDatasetConfiguration(datasetId, optionalDatasetConfigOverride, options);
   }
-  return openTracingViewAndScreenshot(page, baseUrl, createdExplorational.id, optionalViewOverride);
+  await openTracingView(page, baseUrl, createdExplorational.id, optionalViewOverride);
+  return screenshotTracingView(page);
+}
+
+export async function screenshotDatasetWithMapping(
+  page: Page,
+  baseUrl: string,
+  datasetId: APIDatasetId,
+  mappingName: string,
+): Promise<Screenshot> {
+  const options = getDefaultRequestOptions(baseUrl);
+  const createdExplorational = await createExplorational(datasetId, "skeleton", false, options);
+  await openTracingView(page, baseUrl, createdExplorational.id);
+  await page.evaluate(
+    `webknossos.apiReady().then(async api => api.data.activateMapping("${mappingName}"))`,
+  );
+  await waitForMappingEnabled(page);
+  return screenshotTracingView(page);
+}
+
+async function waitForMappingEnabled(page: Page) {
+  let isMappingEnabled;
+  while (!isMappingEnabled) {
+    await page.waitFor(5000);
+    isMappingEnabled = await page.evaluate(
+      "webknossos.apiReady().then(async api => api.data.isMappingEnabled())",
+    );
+  }
 }
 
 async function waitForTracingViewLoad(page: Page) {
   let inputCatchers;
   while (inputCatchers == null || inputCatchers.length < 4) {
-    inputCatchers = await page.$(".inputcatcher");
     await page.waitFor(500);
+    inputCatchers = await page.$(".inputcatcher");
   }
 }
 
@@ -68,20 +95,22 @@ async function waitForRenderingFinish(page: Page) {
   }
 }
 
-async function openTracingViewAndScreenshot(
+async function openTracingView(
   page: Page,
   baseUrl: string,
   annotationId: string,
   optionalViewOverride: ?string,
-): Promise<Screenshot> {
+) {
   const urlSlug = optionalViewOverride != null ? `#${optionalViewOverride}` : "";
   await page.goto(urljoin(baseUrl, `/annotations/Explorational/${annotationId}${urlSlug}`), {
     timeout: 0,
   });
 
   await waitForTracingViewLoad(page);
   await waitForRenderingFinish(page);
+}
 
+async function screenshotTracingView(page: Page): Promise<Screenshot> {
   // Take screenshots of the other rendered planes
   const PLANE_IDS = [
     "#inputcatcher_TDView",