How to decode a COCO RLE binary mask to an image in javascript?

Question

This is an example of a COCO RLE mask - https://pastebin.com/ZhE2en4C

It's an output from a YOLOv8 validation run, taken from the generated predictions.json file.

I'm trying to decode this string in JavaScript and render it on a canvas. The encoded string is valid, because in python I can do this:

from pycocotools import mask as coco_mask
from PIL import Image

example_prediction = {
    "image_id": "102_jpg",
    "category_id": 0,
    "bbox": [153.106, 281.433, 302.518, 130.737],
    "score": 0.8483,
    "segmentation": {
      "size": [640, 640],
      "counts": "<RLE string here>"
    }
  }

def rle_to_bitmap(rle):
  bitmap = coco_mask.decode(rle)
  return bitmap

def show_bitmap(bitmap):
  img = Image.fromarray(bitmap.astype(np.uint8) * 255, mode='L')
  img.show()
  input("Press Enter to continue...")
  img.close()
    

mask_bitmap = rle_to_bitmap(example_prediction["segmentation"])
show_bitmap(mask_bitmap)

And I can see the decoded mask.

Is there a library I can use to decode the same string in JavaScript and convert it to an Image? I tried digging into the source code of pycocotools, but I couldn't do it.

Answer 1

You can draw the mask on a canvas and then export the image if you need.

For the actual drawing you can use two approaches:

1. Decode RLE into binary mask (2d matrix or flattened) and then paint pixels according to that mask
2. Draw mask directly from RLE string on a virtual canvas and then rotate it by 90deg and flip horizontally

Here's the example of both:

// Styling and scaling just for demo
let wrapper = document.createElement("div")
wrapper.style.cssText = `
  transform-origin: left top;
  transform: scale(8);
`
document.body.style.cssText = `
  background-color: #121212;
  margin: 0;
  overflow: hidden;
`
document.body.appendChild(wrapper)

// Helpers
function createCanvas(width, height) {
  let canvas = document.createElement("canvas")

  canvas.style.cssText = `
    border: 1px solid white;
    display: block;
    float: left;
    image-rendering: pixelated;
  `
  canvas.height = height
  canvas.width = width

  // Comment this line if you need only image sources
  wrapper.appendChild(canvas)

  return canvas
}

function randomColorRGBA() {
  return [
        Math.round(Math.random() * 255),
        Math.round(Math.random() * 255),
        Math.round(Math.random() * 255),
        255
      ]
}

// Fast array flattening (faster than Array.proto.flat())
function flatten(arr) {
  const flattened = []

  !(function flat(arr) {
    arr.forEach((el) => {
      if (Array.isArray(el)) flat(el)
      else flattened.push(el)
    })
  })(arr)

  return flattened
}

// Decode from RLE to Binary Mask
// (pass false to flat argument if you need 2d matrix output)
function decodeCocoRLE([rows, cols], counts, flat = true) {
  let pixelPosition = 0,
      binaryMask
  
  if (flat) {
    binaryMask = Array(rows * cols).fill(0)
  } else {
    binaryMask = Array.from({length: rows}, (_) => Array(cols).fill(0))
  }

  for (let i = 0, rleLength = counts.length; i < rleLength; i += 2) {
    let zeros = counts[i],
        ones = counts[i + 1] ?? 0

    pixelPosition += zeros 

    while (ones > 0) {
      const rowIndex = pixelPosition % rows,
            colIndex = (pixelPosition - rowIndex) / rows

      if (flat) {
        const arrayIndex = rowIndex * cols + colIndex
        binaryMask[arrayIndex] = 1
      } else {
        binaryMask[rowIndex][colIndex] = 1
      }

      pixelPosition++
      ones--
    }
  }

  if (!flat) {
    console.log("Result matrix:")
    binaryMask.forEach((row, i) => console.log(row.join(" "), `- row ${i}`))
  }

  return binaryMask
}

// 1. Draw from binary mask
function drawFromBinaryMask({size, counts}) {
  let fillColor = randomColorRGBA(),
      height = size[0],
      width = size[1]

  let canvas = createCanvas(width, height),
      canvasCtx = canvas.getContext("2d"),
      imgData = canvasCtx.getImageData(0, 0, width, height),
      pixelData = imgData.data

  // If you need matrix output (flat = false)
  // let maskFlattened = flatten(decodeCocoRLE(size, counts, false)),
  //     maskLength = maskFlattened.length;
  
  // If not - it's better to use faster approach
  let maskFlattened = decodeCocoRLE(size, counts),
      maskLength = maskFlattened.length;

  for(let i = 0; i < maskLength; i++) {
    if (maskFlattened[i] === 1) {
      let pixelPosition = i * 4

      pixelData[pixelPosition] = fillColor[0]
      pixelData[pixelPosition + 1] = fillColor[1]
      pixelData[pixelPosition + 2] = fillColor[2]
      pixelData[pixelPosition + 3] = fillColor[3]
    }
  }

  canvasCtx.putImageData(imgData, 0, 0)

  // If needed you can return data:image/png 
  // to use it as an image.src
  return canvas.toDataURL()
}

// 2. Draw using virtual canvas
function drawDirectlyFromRle({size: [rows, cols], counts}) {
  let fillColor = randomColorRGBA(),
      isOnesInterval = false,
      start = 0,
      end = 0

  let realCanvas = createCanvas(cols, rows),
      realCtx = realCanvas.getContext("2d")

  let virtualCanvas = new OffscreenCanvas(rows, cols),
      virtualCtx = virtualCanvas.getContext("2d"),
      imgData = virtualCtx.getImageData(0, 0, rows, cols),
      pixelData = imgData.data

  counts.forEach((interval) => {
    end = start + interval * 4
    if (isOnesInterval) {
      for (let i = start; i < end; i += 4) {
        pixelData[i] = fillColor[0]
        pixelData[i + 1] = fillColor[1]
        pixelData[i + 2] = fillColor[2]
        pixelData[i + 3] = fillColor[3]
      }
    }
    start = end
    isOnesInterval = !isOnesInterval
  })

  virtualCtx.putImageData(imgData, 0, 0)

  realCtx.save()
  realCtx.scale(-1, 1)
  realCtx.rotate(90*Math.PI/180)
  realCtx.drawImage(virtualCanvas, 0, 0)
  realCtx.restore()

  // If needed you can return data:image/png 
  // to use it as an image.src
  return realCanvas.toDataURL()
}

// Test RLE
const exampleCocoRLE = {
        counts: [15, 1, 9, 1, 3, 3, 2, 1, 8, 1, 8, 1, 3, 3, 2, 1, 8, 1, 7, 1, 11],
        size: [9, 10]
    }

// Draw on canvas
let imageSrc1 = drawFromBinaryMask(exampleCocoRLE),
    imageSrc2 = drawDirectlyFromRle(exampleCocoRLE)

console.log("Canvas 1 image (from binary):\n", imageSrc1)
console.log("Canvas 2 image (from virtual):\n", imageSrc2)

// Example of src usage
let image1 = document.createElement("img"),
    image2 = document.createElement("img"),
    imageStyle = `
      display: block;
      float: left;
      border: 1px solid lime;
      image-rendering: pixelated;
    `

// demo styling
image1.style.cssText = imageStyle
image2.style.cssText = imageStyle

image1.onload = () => {
  wrapper.appendChild(image1)
}
image2.onload = () => {
  wrapper.appendChild(image2)
}

image1.src = imageSrc1
image2.src = imageSrc2