浏览器中ONNX分割模型输出的后处理（掩码转换与尺寸调整）技术求助

嘿，刚好之前在浏览器端折腾过ONNX分割模型的后处理，给你分享下我当时的实现思路和代码片段，应该能帮你解决问题！

一、把模型输出转成可用掩码 #

首先得看你的模型是二分类还是多分类，两种情况的处理逻辑略有不同：

情况1：二分类分割模型（比如前景/背景分割） #

这类模型的输出通常是[1, 1, maskH, maskW]形状的张量，每个像素是一个概率值（0到1之间）。我们可以设置一个阈值，把概率大于阈值的像素标记为前景，否则为背景：

// 输入：ONNX Runtime返回的输出张量
async function convertBinaryOutputToMask(outputTensor) {
  // 提取张量的形状和原始数据
  const [_, __, maskH, maskW] = outputTensor.dims;
  const outputData = outputTensor.data;
  
  // 初始化掩码数组，用Uint8Array存储更节省内存
  const maskData = new Uint8Array(maskH * maskW);
  
  // 遍历每个像素，根据阈值生成掩码
  const threshold = 0.5; // 可以根据你的模型调整这个值
  for (let i = 0; i < maskH * maskW; i++) {
    // 把前景设为255（方便后续可视化），背景设为0
    maskData[i] = outputData[i] > threshold ? 255 : 0;
  }
  
  return { maskData, maskH, maskW };
}

情况2：多分类分割模型 #

这类模型输出是[1, numClasses, maskH, maskW]形状的张量，每个像素对应多个类别的概率值。我们需要对每个像素取概率最大的类别作为该像素的分类结果：

async function convertMultiClassOutputToMask(outputTensor) {
  const [_, numClasses, maskH, maskW] = outputTensor.dims;
  const outputData = outputTensor.data;
  const maskData = new Uint8Array(maskH * maskW);
  
  // 遍历每个像素，找到概率最高的类别
  for (let y = 0; y < maskH; y++) {
    for (let x = 0; x < maskW; x++) {
      let maxProb = -Infinity;
      let targetClass = 0;
      
      // 遍历当前像素的所有类别概率
      for (let c = 0; c < numClasses; c++) {
        const idx = c * maskH * maskW + y * maskW + x;
        if (outputData[idx] > maxProb) {
          maxProb = outputData[idx];
          targetClass = c;
        }
      }
      
      // 把类别索引转成可视化的数值（这里给每个类别分配不同的亮度值）
      maskData[y * maskW + x] = targetClass * 50;
    }
  }
  
  return { maskData, maskH, maskW };
}

二、掩码的尺寸调整与可视化 #

通常模型输出的掩码尺寸会比输入图片小（比如输入是640x480，输出是320x240），这时候我们需要把掩码 resize 到和原图一致的尺寸，再叠加到原图上可视化。

我当时用的是OffscreenCanvas来做缩放，效率比普通Canvas高，不会阻塞主线程：

// 调整掩码尺寸到目标大小（比如原图的宽高）
function resizeMask(maskData, originalMaskW, originalMaskH, targetW, targetH) {
  // 创建原始尺寸的离屏Canvas，把掩码画上去
  const originalCanvas = new OffscreenCanvas(originalMaskW, originalMaskH);
  const originalCtx = originalCanvas.getContext('2d');
  const originalImageData = originalCtx.createImageData(originalMaskW, originalMaskH);
  
  // 把掩码数据填充到ImageData中
  for (let i = 0; i < maskData.length; i++) {
    // 单通道掩码，把R/G/B都设为掩码值，Alpha通道设为255
    originalImageData.data[i * 4] = maskData[i];
    originalImageData.data[i * 4 + 1] = maskData[i];
    originalImageData.data[i * 4 + 2] = maskData[i];
    originalImageData.data[i * 4 + 3] = 255;
  }
  originalCtx.putImageData(originalImageData, 0, 0);
  
  // 创建目标尺寸的离屏Canvas，进行缩放绘制
  const targetCanvas = new OffscreenCanvas(targetW, targetH);
  const targetCtx = targetCanvas.getContext('2d');
  targetCtx.drawImage(originalCanvas, 0, 0, targetW, targetH);
  
  // 提取缩放后的掩码数据
  const resizedImageData = targetCtx.getImageData(0, 0, targetW, targetH);
  const resizedMaskData = new Uint8Array(targetW * targetH);
  for (let i = 0; i < targetW * targetH; i++) {
    resizedMaskData[i] = resizedImageData.data[i * 4];
  }
  
  return resizedMaskData;
}

最后是可视化：把掩码叠加到原图上 #

你可以在页面上放一个Canvas，把原图和半透明的掩码画在一起：

// 输入：页面上的原图元素、掩码数据、掩码原始尺寸
async function renderSegmentationResult(originalImgEl, maskData, maskW, maskH) {
  const canvas = document.getElementById('seg-result-canvas');
  canvas.width = originalImgEl.width;
  canvas.height = originalImgEl.height;
  const ctx = canvas.getContext('2d');
  
  // 先画原图
  ctx.drawImage(originalImgEl, 0, 0);
  
  // 调整掩码到原图尺寸
  const resizedMask = resizeMask(maskData, maskW, maskH, canvas.width, canvas.height);
  
  // 创建掩码的ImageData，设置半透明效果
  const maskImageData = ctx.createImageData(canvas.width, canvas.height);
  for (let i = 0; i < resizedMask.length; i++) {
    if (resizedMask[i] > 0) {
      // 这里用红色半透明作为前景标记，你可以改成其他颜色
      maskImageData.data[i * 4] = 255; // R
      maskImageData.data[i * 4 + 1] = 0; // G
      maskImageData.data[i * 4 + 2] = 0; // B
      maskImageData.data[i * 4 + 3] = 80; // Alpha，值越小越透明
    } else {
      maskImageData.data[i * 4 + 3] = 0; // 背景全透明
    }
  }
  
  // 把掩码画到Canvas上，叠加在原图上方
  ctx.putImageData(maskImageData, 0, 0);
}

一些小提示 #

• 阈值调整：二分类的阈值可以根据你的模型实际输出调整，比如如果模型输出的前景概率普遍偏低，可以把阈值设为0.3试试
• 多分类颜色映射：如果是多分类，你可以给每个类别定义不同的RGB颜色，比如用一个数组存储[[255,0,0], [0,255,0], [0,0,255]]这样的颜色值，然后根据classIdx取对应的颜色
• 性能优化：如果你的模型输出尺寸很大，尽量用TypedArray（比如Uint8Array）存储数据，比普通数组效率高很多

备注：内容来源于stack exchange，提问作者Krishna Kurdekar