How to decode COCO RLE binary mask to image in JavaScript?

Question

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

This is the output of the YOLOv8 validation run, taken from the generated Predictions.json file.

I'm trying to decode the string in JavaScript and render it on the 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)

I can see the decoded mask.

Is there a library that can be used 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.

Answer 1

You can draw the mask on the canvas and then export the image if needed.

For actual drawing, you can use two methods:

1. Decode the RLE into a binary mask (2D matrix or flattened matrix) and draw pixels based on that mask
2. Draw the mask directly from the RLE string on the virtual canvas, then rotate it 90 degrees and flip it horizontally

Here are examples 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