Function toSegmentationImage

deeplab/src/utils.ts:184–223

(
    colormap: Color[],
    labelNames: string[],
    rawSegmentationMap: tf.Tensor2D,
    canvas?: HTMLCanvasElement,
    )

 *   fed into `ImageData` and mapped to a canvas.
 */
export async function toSegmentationImage(
    colormap: Color[],
    labelNames: string[],
    rawSegmentationMap: tf.Tensor2D,
    canvas?: HTMLCanvasElement,
    ): Promise<SegmentationData> {
  if (colormap.length < labelNames.length) {
    throw new Error(
        'The colormap must be expansive enough to encode each label. ' +
        `Aborting, since the given colormap has length ${colormap.length}, ` +
        `but there are ${labelNames.length} labels.`);
  }
  const [height, width] = rawSegmentationMap.shape;
  const segmentationImageBuffer = tf.buffer([height, width, 3], 'int32');
  const mapData = await rawSegmentationMap.array();
  const labels = new Set<Label>();
  for (let columnIndex = 0; columnIndex < height; ++columnIndex) {
    for (let rowIndex = 0; rowIndex < width; ++rowIndex) {
      const label: Label = mapData[columnIndex][rowIndex];
      labels.add(label);
      segmentationImageBuffer.set(colormap[label][0], columnIndex, rowIndex, 0);
      segmentationImageBuffer.set(colormap[label][1], columnIndex, rowIndex, 1);
      segmentationImageBuffer.set(colormap[label][2], columnIndex, rowIndex, 2);
    }
  }

  const segmentationImageTensor =
      segmentationImageBuffer.toTensor() as tf.Tensor3D;

  const segmentationMap =
      await tf.browser.toPixels(segmentationImageTensor, canvas);

  tf.dispose(segmentationImageTensor);

  const legend: Legend = {};
  for (const label of Array.from(labels)) {
    legend[labelNames[label]] = colormap[label];
  }
  return {legend, segmentationMap};
}