如何在JavaScript中将COCO RLE二进制掩码解码为图像?

bq3bfh9z  于 2023-08-02  发布在  Java
关注(0)|答案(1)|浏览(158)

这是COCO RLE掩码的示例-https://pastebin.com/ZhE2en4C
它是YOLOv 8验证运行的输出,取自生成的predictions.json文件。
我正在尝试用JavaScript解码这个字符串并将其呈现在画布上。编码的字符串是有效的,因为在python中我可以这样做:

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)

字符串
我能看到解码的面具。
有没有一个库可以用来在JavaScript中解码相同的字符串并将其转换为Image?我试着深入挖掘pycocotools的源代码,但我做不到。

3pvhb19x

3pvhb19x1#

您可以在画布上绘制蒙版,然后根据需要导出图像。
对于实际绘图,可以使用两种方法:
1.将RLE解码为二进制掩码,然后根据该掩码绘制像素
1.直接从虚拟画布上的RLE字符串绘制蒙版,然后将其旋转90度并水平翻转
以下是两者的示例(以全页模式打开以查看缩放结果):

// Styling and scaling just for demo
let wrapper = document.createElement("div")
wrapper.style.transformOrigin = "left top"
wrapper.style.transform = "scale(10)"
document.body.style.backgroundColor = '#121212'
document.body.style.overflow = 'hidden'
document.body.appendChild(wrapper)

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

  canvas.style.border = "1px solid white"
  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
      ]
}

// Decode from RLE to Binary Mask Matrix
// (slow approach, just is case you need matrix representation)
function decodeCocoRLE([rows, cols], counts) {
  // Init matrix filled with 'zeros', size [cols x rows]
  // Matrix actually is a two dimensional array
  // with 'rows' nested arrays, 'cols' elements each
  let binaryMask = Array.from({length: rows}, (_, i) => Array(cols).fill(0)),
      rleLength = counts.length,
      processedPixelsCount = 0;

  // All we need to do is parse the RLE 
  // and populate our matrix with 'ones'
  for (let i = 0; i < rleLength; i += 2) {
    let zeros = counts[i],
        ones = counts[i + 1] ?? 0,
        pixelPosition = processedPixelsCount + zeros

    processedPixelsCount += zeros + ones

    while (ones > 0) {
      // Calculate the pixel position in the matrix
      const rowIndex = pixelPosition % rows;
      const colIndex = (pixelPosition - rowIndex) / rows;

      binaryMask[rowIndex][colIndex] = 1
      pixelPosition++
      ones--
    }
  }

  console.log("Result matrix:")
  binaryMask.forEach((row) => console.log(row.join(" ")))
  
  return binaryMask
}

// Decode from RLE to Flattened Binary Mask
// (much faster approach)
function decodeCocoRLEtoFlat([rows, cols], counts) {
  let binaryMaskFlattened = Array(rows * cols).fill(0),
      processedPixelsCount = 0;

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

    processedPixelsCount += zeros + ones

    while (ones > 0) {
      const rowIndex = pixelPosition % rows
      const colIndex = (pixelPosition - rowIndex) / rows
      const arrayIndex = rowIndex * cols + colIndex

      binaryMaskFlattened[arrayIndex] = 1
      pixelPosition++
      ones--
    }
  }

  return binaryMaskFlattened
}

// 1. Draw from 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
  // let maskFlattened = decodeCocoRLE(size, counts).flat(),
  //     maskLength = maskFlattened.length;
  
  // If not - it's better to use faster approach
  let maskFlattened = decodeCocoRLEtoFlat(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 = document.createElement("canvas"),
      virtualCtx = virtualCanvas.getContext("2d")

  virtualCanvas.width = rows
  virtualCanvas.height = cols

  let 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: [6, 1, 40, 4, 5, 4, 5, 4, 21], size: [9, 10]}

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

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

// Example of src usage
let image1 = document.createElement("img"),
    image2 = document.createElement("img")

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

image1.src = imageSrc1
image2.src = imageSrc2

字符串

相关问题