PaintFilter

Source

PaintFilter.worker.js

import registerWebworker from 'webworker-promise/lib/register';

import { SlicingMode } from 'vtk.js/Sources/Rendering/Core/ImageMapper/Constants';
import { vec3 } from 'gl-matrix';

const globals = {
  // single-component labelmap
  buffer: null,
  dimensions: [0, 0, 0],
  prevPoint: null,
  slicingMode: null, // 2D or 3D painting
};

// --------------------------------------------------------------------------

function handlePaintRectangle({ point1, point2 }) {
  const [x1, y1, z1] = point1;
  const [x2, y2, z2] = point2;

  const xstart = Math.max(Math.min(x1, x2), 0);
  const xend = Math.min(Math.max(x1, x2), globals.dimensions[0] - 1);
  if (xstart <= xend) {
    const ystart = Math.max(Math.min(y1, y2), 0);
    const yend = Math.min(Math.max(y1, y2), globals.dimensions[1] - 1);
    const zstart = Math.max(Math.min(z1, z2), 0);
    const zend = Math.min(Math.max(z1, z2), globals.dimensions[2] - 1);

    const jStride = globals.dimensions[0];
    const kStride = globals.dimensions[0] * globals.dimensions[1];

    for (let k = zstart; k <= zend; k++) {
      for (let j = ystart; j <= yend; j++) {
        const index = j * jStride + k * kStride;
        globals.buffer.fill(1, index + xstart, index + xend + 1);
      }
    }
  }
}

// --------------------------------------------------------------------------
// center and scale3 are in IJK coordinates
function handlePaintEllipse({ center, scale3 }) {
  const radius3 = [...scale3];
  const indexCenter = center.map((val) => Math.round(val));

  let sliceAxis = -1;
  if (globals.slicingMode != null && globals.slicingMode !== SlicingMode.NONE) {
    sliceAxis = globals.slicingMode % 3;
  }

  const yStride = globals.dimensions[0];
  const zStride = globals.dimensions[0] * globals.dimensions[1];

  let [xmin, ymin, zmin] = indexCenter;
  let [xmax, ymax, zmax] = indexCenter;

  if (sliceAxis !== 2) {
    zmin = Math.round(Math.max(indexCenter[2] - radius3[2], 0));
    zmax = Math.round(
      Math.min(indexCenter[2] + radius3[2], globals.dimensions[2] - 1)
    );
  }

  for (let z = zmin; z <= zmax; z++) {
    let dz = 0;
    if (sliceAxis !== 2) {
      dz = (indexCenter[2] - z) / radius3[2];
    }

    const dzSquared = dz * dz;

    if (dzSquared <= 1) {
      const ay = radius3[1] * Math.sqrt(1 - dzSquared);
      if (sliceAxis !== 1) {
        ymin = Math.round(Math.max(indexCenter[1] - ay, 0));
        ymax = Math.round(
          Math.min(indexCenter[1] + ay, globals.dimensions[1] - 1)
        );
      }

      for (let y = ymin; y <= ymax; y++) {
        let dy = 0;
        if (sliceAxis !== 1) {
          dy = (indexCenter[1] - y) / radius3[1];
        }
        const dySquared = dy * dy;
        if (dySquared + dzSquared <= 1) {
          if (sliceAxis !== 0) {
            const ax = radius3[0] * Math.sqrt(1 - dySquared - dzSquared);
            xmin = Math.round(Math.max(indexCenter[0] - ax, 0));
            xmax = Math.round(
              Math.min(indexCenter[0] + ax, globals.dimensions[0] - 1)
            );
          }
          if (xmin <= xmax) {
            const index = y * yStride + z * zStride;
            globals.buffer.fill(1, index + xmin, index + xmax + 1);
          }
        }
      }
    }
  }
}

// --------------------------------------------------------------------------

function handlePaint({ point, radius }) {
  if (!globals.prevPoint) {
    globals.prevPoint = point;
  }

  // DDA params
  const delta = [
    point[0] - globals.prevPoint[0],
    point[1] - globals.prevPoint[1],
    point[2] - globals.prevPoint[2],
  ];
  const inc = [1, 1, 1];
  for (let i = 0; i < 3; i++) {
    if (delta[i] < 0) {
      delta[i] = -delta[i];
      inc[i] = -1;
    }
  }
  const step = Math.max(...delta);

  // DDA
  const thresh = [step, step, step];
  const pt = [...globals.prevPoint];
  for (let s = 0; s <= step; s++) {
    handlePaintEllipse({ center: pt, scale3: radius });

    for (let ii = 0; ii < 3; ii++) {
      thresh[ii] -= delta[ii];
      if (thresh[ii] <= 0) {
        thresh[ii] += step;
        pt[ii] += inc[ii];
      }
    }
  }

  globals.prevPoint = point;
}

// --------------------------------------------------------------------------

function handlePaintTriangles({ triangleList }) {
  // debugger;

  const triangleCount = Math.floor(triangleList.length / 9);

  for (let i = 0; i < triangleCount; i++) {
    const point0 = triangleList.subarray(9 * i + 0, 9 * i + 3);
    const point1 = triangleList.subarray(9 * i + 3, 9 * i + 6);
    const point2 = triangleList.subarray(9 * i + 6, 9 * i + 9);

    const v1 = [0, 0, 0];
    const v2 = [0, 0, 0];

    vec3.subtract(v1, point1, point0);
    vec3.subtract(v2, point2, point0);

    const step1 = [0, 0, 0];
    const numStep1 =
      2 * Math.max(Math.abs(v1[0]), Math.abs(v1[1]), Math.abs(v1[2]));
    vec3.scale(step1, v1, 1 / numStep1);

    const step2 = [0, 0, 0];
    const numStep2 =
      2 * Math.max(Math.abs(v2[0]), Math.abs(v2[1]), Math.abs(v2[2]));
    vec3.scale(step2, v2, 1 / numStep2);

    const jStride = globals.dimensions[0];
    const kStride = globals.dimensions[0] * globals.dimensions[1];

    for (let u = 0; u <= numStep1 + 1; u++) {
      const maxV = numStep2 - u * (numStep2 / numStep1);
      for (let v = 0; v <= maxV + 1; v++) {
        const point = [...point0];
        vec3.scaleAndAdd(point, point, step1, u);
        vec3.scaleAndAdd(point, point, step2, v);

        point[0] = Math.round(point[0]);
        point[1] = Math.round(point[1]);
        point[2] = Math.round(point[2]);

        if (
          point[0] >= 0 &&
          point[0] < globals.dimensions[0] &&
          point[1] >= 0 &&
          point[1] < globals.dimensions[1] &&
          point[2] >= 0 &&
          point[2] < globals.dimensions[2]
        ) {
          globals.buffer[
            point[0] + jStride * point[1] + kStride * point[2]
          ] = 1;
        }
      }
    }
  }
}

// --------------------------------------------------------------------------

registerWebworker()
  .operation('start', ({ bufferType, dimensions, slicingMode }) => {
    if (!globals.buffer) {
      const bufferSize = dimensions[0] * dimensions[1] * dimensions[2];
      /* eslint-disable-next-line */
      globals.buffer = new self[bufferType](bufferSize);
      globals.dimensions = dimensions;
      globals.prevPoint = null;
      globals.slicingMode = slicingMode;
    }
  })
  .operation('paint', handlePaint)
  .operation('paintRectangle', handlePaintRectangle)
  .operation('paintEllipse', handlePaintEllipse)
  .operation('paintTriangles', handlePaintTriangles)
  .operation('end', () => {
    const response = new registerWebworker.TransferableResponse(
      globals.buffer.buffer,
      [globals.buffer.buffer]
    );
    globals.buffer = null;
    return response;
  });

index.js

import { vec3 } from 'gl-matrix';
import WebworkerPromise from 'webworker-promise';

import macro from 'vtk.js/Sources/macros';
import vtkImageData from 'vtk.js/Sources/Common/DataModel/ImageData';
import vtkDataArray from 'vtk.js/Sources/Common/Core/DataArray';
import vtkPolygon from 'vtk.js/Sources/Common/DataModel/Polygon';

import PaintFilterWorker from 'vtk.js/Sources/Filters/General/PaintFilter/PaintFilter.worker';

const { vtkErrorMacro } = macro;

// ----------------------------------------------------------------------------
// vtkPaintFilter methods
// ----------------------------------------------------------------------------

function vtkPaintFilter(publicAPI, model) {
  // Set our className
  model.classHierarchy.push('vtkPaintFilter');

  let worker = null;
  let workerPromise = null;
  const history = {};

  // --------------------------------------------------------------------------

  function resetHistory() {
    history.index = -1;
    history.snapshots = [];
    history.labels = [];
  }

  function pushToHistory(snapshot, label) {
    // Clear any "redo" info
    const spliceIndex = history.index + 1;
    const spliceLength = history.snapshots.length - history.index;
    history.snapshots.splice(spliceIndex, spliceLength);
    history.labels.splice(spliceIndex, spliceLength);

    // Push new snapshot
    history.snapshots.push(snapshot);
    history.labels.push(label);
    history.index++;
  }

  // --------------------------------------------------------------------------

  publicAPI.startStroke = () => {
    if (model.labelMap) {
      if (!workerPromise) {
        worker = new PaintFilterWorker();
        workerPromise = new WebworkerPromise(worker);
      }

      workerPromise.exec('start', {
        bufferType: 'Uint8Array',
        dimensions: model.labelMap.getDimensions(),
        slicingMode: model.slicingMode,
      });
    }
  };

  // --------------------------------------------------------------------------

  publicAPI.endStroke = () => {
    let endStrokePromise;

    if (workerPromise) {
      endStrokePromise = workerPromise.exec('end');

      endStrokePromise.then((strokeBuffer) => {
        publicAPI.applyBinaryMask(strokeBuffer);
        worker.terminate();
        worker = null;
        workerPromise = null;
      });
    }
    return endStrokePromise;
  };

  publicAPI.applyBinaryMask = (maskBuffer) => {
    const scalars = model.labelMap.getPointData().getScalars();
    const data = scalars.getData();
    const maskLabelMap = new Uint8Array(maskBuffer);

    let diffCount = 0;
    for (let i = 0; i < maskLabelMap.length; i++) {
      // maskLabelMap is a binary mask
      diffCount += maskLabelMap[i];
    }

    // Format: [ [index, oldLabel], ...]
    // I could use an ArrayBuffer, which would place limits
    // on the values of index/old, but will be more efficient.
    const snapshot = new Array(diffCount);
    const label = model.label;

    let diffIdx = 0;
    if (model.voxelFunc) {
      const bgScalars = model.backgroundImage.getPointData().getScalars();
      const voxel = [];
      for (let i = 0; i < maskLabelMap.length; i++) {
        if (maskLabelMap[i]) {
          bgScalars.getTuple(i, voxel);
          // might not fill up snapshot
          if (model.voxelFunc(voxel, i, label)) {
            snapshot[diffIdx++] = [i, data[i]];
            data[i] = label;
          }
        }
      }
    } else {
      for (let i = 0; i < maskLabelMap.length; i++) {
        if (maskLabelMap[i]) {
          if (data[i] !== label) {
            snapshot[diffIdx++] = [i, data[i]];
            data[i] = label;
          }
        }
      }
    }
    pushToHistory(snapshot, label);

    scalars.setData(data);
    scalars.modified();
    model.labelMap.modified();
    publicAPI.modified();
  };

  // --------------------------------------------------------------------------

  publicAPI.addPoint = (point) => {
    if (workerPromise) {
      const worldPt = [point[0], point[1], point[2]];
      const indexPt = [0, 0, 0];
      vec3.transformMat4(indexPt, worldPt, model.maskWorldToIndex);
      indexPt[0] = Math.round(indexPt[0]);
      indexPt[1] = Math.round(indexPt[1]);
      indexPt[2] = Math.round(indexPt[2]);

      const spacing = model.labelMap.getSpacing();
      const radius = spacing.map((s) => model.radius / s);

      workerPromise.exec('paint', { point: indexPt, radius });
    }
  };

  // --------------------------------------------------------------------------

  publicAPI.paintRectangle = (point1, point2) => {
    if (workerPromise) {
      const index1 = [0, 0, 0];
      const index2 = [0, 0, 0];
      vec3.transformMat4(index1, point1, model.maskWorldToIndex);
      vec3.transformMat4(index2, point2, model.maskWorldToIndex);
      index1[0] = Math.round(index1[0]);
      index1[1] = Math.round(index1[1]);
      index1[2] = Math.round(index1[2]);
      index2[0] = Math.round(index2[0]);
      index2[1] = Math.round(index2[1]);
      index2[2] = Math.round(index2[2]);
      workerPromise.exec('paintRectangle', {
        point1: index1,
        point2: index2,
      });
    }
  };

  // --------------------------------------------------------------------------

  publicAPI.paintEllipse = (center, scale3) => {
    if (workerPromise) {
      const realCenter = [0, 0, 0];
      const origin = [0, 0, 0];
      let realScale3 = [0, 0, 0];
      vec3.transformMat4(realCenter, center, model.maskWorldToIndex);
      vec3.transformMat4(origin, origin, model.maskWorldToIndex);
      vec3.transformMat4(realScale3, scale3, model.maskWorldToIndex);
      vec3.subtract(realScale3, realScale3, origin);
      realScale3 = realScale3.map((s) => (s === 0 ? 0.25 : Math.abs(s)));
      workerPromise.exec('paintEllipse', {
        center: realCenter,
        scale3: realScale3,
      });
    }
  };

  // --------------------------------------------------------------------------

  publicAPI.canUndo = () => history.index > -1;

  // --------------------------------------------------------------------------

  publicAPI.paintPolygon = (pointList) => {
    if (workerPromise && pointList.length > 0) {
      const polygon = vtkPolygon.newInstance();
      const poly = [];
      for (let i = 0; i < pointList.length / 3; i++) {
        poly.push([
          pointList[3 * i + 0],
          pointList[3 * i + 1],
          pointList[3 * i + 2],
        ]);
      }
      polygon.setPoints(poly);

      if (!polygon.triangulate()) {
        console.log('triangulation failed!');
      }

      const points = polygon.getPointArray();
      const triangleList = new Float32Array(points.length);
      const numPoints = Math.floor(triangleList.length / 3);
      for (let i = 0; i < numPoints; i++) {
        const point = points.slice(3 * i, 3 * i + 3);
        const voxel = triangleList.subarray(3 * i, 3 * i + 3);
        vec3.transformMat4(voxel, point, model.maskWorldToIndex);
      }

      workerPromise.exec('paintTriangles', {
        triangleList,
      });
    }
  };

  // --------------------------------------------------------------------------

  publicAPI.applyLabelMap = (labelMap) => {
    const currentMapData = model.labelMap.getPointData().getScalars().getData();

    const newMapData = labelMap.getPointData().getScalars().getData();

    // Compute snapshot
    const snapshot = [];
    for (let i = 0; i < newMapData.length; ++i) {
      if (currentMapData[i] !== newMapData[i]) {
        snapshot.push([i, currentMapData[i]]);
      }
    }

    pushToHistory(snapshot, model.label);
    model.labelMap = labelMap;
    publicAPI.modified();
  };

  // --------------------------------------------------------------------------

  publicAPI.undo = () => {
    if (history.index > -1) {
      const scalars = model.labelMap.getPointData().getScalars();
      const data = scalars.getData();

      const snapshot = history.snapshots[history.index];
      for (let i = 0; i < snapshot.length; i++) {
        if (!snapshot[i]) {
          break;
        }

        const [index, oldLabel] = snapshot[i];
        data[index] = oldLabel;
      }

      history.index--;

      scalars.setData(data);
      scalars.modified();
      model.labelMap.modified();
      publicAPI.modified();
    }
  };

  // --------------------------------------------------------------------------

  publicAPI.canRedo = () => history.index < history.labels.length - 1;

  // --------------------------------------------------------------------------

  publicAPI.redo = () => {
    if (history.index < history.labels.length - 1) {
      const scalars = model.labelMap.getPointData().getScalars();
      const data = scalars.getData();

      const redoLabel = history.labels[history.index + 1];
      const snapshot = history.snapshots[history.index + 1];

      for (let i = 0; i < snapshot.length; i++) {
        if (!snapshot[i]) {
          break;
        }

        const [index] = snapshot[i];
        data[index] = redoLabel;
      }

      history.index++;

      scalars.setData(data);
      scalars.modified();
      model.labelMap.modified();
      publicAPI.modified();
    }
  };

  // --------------------------------------------------------------------------

  const superSetLabelMap = publicAPI.setLabelMap;
  publicAPI.setLabelMap = (lm) => {
    if (superSetLabelMap(lm)) {
      model.maskWorldToIndex = model.labelMap.getWorldToIndex();
      resetHistory();
      return true;
    }
    return false;
  };

  // --------------------------------------------------------------------------

  publicAPI.requestData = (inData, outData) => {
    if (!model.backgroundImage) {
      vtkErrorMacro('No background image');
      return;
    }

    if (!model.backgroundImage.getPointData().getScalars()) {
      vtkErrorMacro('Background image has no scalars');
      return;
    }

    if (!model.labelMap) {
      // clone background image properties
      const labelMap = vtkImageData.newInstance(
        model.backgroundImage.get('spacing', 'origin', 'direction')
      );
      labelMap.setDimensions(model.backgroundImage.getDimensions());
      labelMap.computeTransforms();

      // right now only support 256 labels
      const values = new Uint8Array(model.backgroundImage.getNumberOfPoints());
      const dataArray = vtkDataArray.newInstance({
        numberOfComponents: 1, // labelmap with single component
        values,
      });
      labelMap.getPointData().setScalars(dataArray);

      publicAPI.setLabelMap(labelMap);
    }

    if (!model.maskWorldToIndex) {
      model.maskWorldToIndex = model.labelMap.getWorldToIndex();
    }

    const scalars = model.labelMap.getPointData().getScalars();

    if (!scalars) {
      vtkErrorMacro('Mask image has no scalars');
      return;
    }

    model.labelMap.modified();

    outData[0] = model.labelMap;
  };

  // --------------------------------------------------------------------------
  // Initialization
  // --------------------------------------------------------------------------

  resetHistory();
}

// ----------------------------------------------------------------------------
// Object factory
// ----------------------------------------------------------------------------

const DEFAULT_VALUES = {
  backgroundImage: null,
  labelMap: null,
  maskWorldToIndex: null,
  voxelFunc: null,
  radius: 1,
  label: 0,
  slicingMode: null,
};

// ----------------------------------------------------------------------------

export function extend(publicAPI, model, initialValues = {}) {
  Object.assign(model, DEFAULT_VALUES, initialValues);

  // Make this a VTK object
  macro.obj(publicAPI, model);

  // Also make it an algorithm with no input and one output
  macro.algo(publicAPI, model, 0, 1);

  macro.setGet(publicAPI, model, [
    'backgroundImage',
    'labelMap',
    'maskWorldToIndex',
    'voxelFunc',
    'label',
    'radius',
    'slicingMode',
  ]);

  // Object specific methods
  vtkPaintFilter(publicAPI, model);
}

// ----------------------------------------------------------------------------

export const newInstance = macro.newInstance(extend, 'vtkPaintFilter');

// ----------------------------------------------------------------------------

export default { newInstance, extend };