process-output.js

const colorMapUrl = '/assets/color-map.json'

// load the color-map file
const response = await fetch(colorMapUrl)
const colorMap = await response.json()

// get image width, height from output
const imageWidth = prediction.shape[2]
const imageHeight = prediction.shape[1]

// create a regular array from the model's output prediction (tensor)
const segMap = Array.from(prediction.dataSync())

// create a new array with the corresponding segmentation colors
const segMapColor = segMap.map(seg => colorMap[seg])

let data = []

// load segmentation colors into format required for images
for (var i = 0; i < segMapColor.length; i++) {
  data.push(segMapColor[i][0]) // red
  data.push(segMapColor[i][1]) // green
  data.push(segMapColor[i][2]) // blue
  data.push(175) // alpha
}

// create an image with the segmentation color data
let imageData = new ImageData(imageWidth, imageHeight)
imageData.data.set(data)