Comment décoder le masque binaire COCO RLE en image en JavaScript ?

Question

Ceci est un exemple de masque COCO RLE - https://pastebin.com/ZhE2en4C

Il s'agit du résultat de l'exécution de validation YOLOv8, extrait du fichier Predictions.json généré.

J'essaie de décoder cette chaîne en JavaScript et de la restituer sur le canevas. La chaîne encodée est valide car en python je peux faire ceci :

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)

Je peux voir le masque décodé.

Existe-t-il une bibliothèque qui peut être utilisée pour décoder la même chaîne en JavaScript et la convertir en Image ? J'ai essayé de fouiller dans le code source de pycocotools mais je n'ai pas pu.

Answer 1

Vous pouvez dessiner le masque sur la toile puis exporter l'image si nécessaire.

Pour le dessin proprement dit, vous pouvez utiliser deux méthodes :

1. Décodez le RLE en un masque binaire (matrice 2D ou matrice aplatie) et dessinez des pixels en fonction de ce masque
2. Dessinez le masque directement à partir de la chaîne RLE sur la toile virtuelle, puis faites-le pivoter de 90 degrés et retournez-le horizontalement

Voici des exemples des deux :

// 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