TensorFlow.js DeepLab模型形状不匹配错误求助

问题重现 #

执行模型时触发形状不匹配错误：

Uncaught (in promise) Error: The shape of dict['ImageTensor'] provided in model.execute(dict) must be [1,-1,-1,3], but was [1,1,1,513,3]

环境信息：TensorFlow.js 2.8.5，通过@tensorflow-models/deeplab加载模型。

问题原因 #

model.segment()方法本身支持直接传入HTMLCanvasElement/HTMLImageElement，并会自动完成张量转换、添加批量维度和预处理操作。手动对图像进行归一化并添加批量维度后再传入该方法，会导致内部再次对已有批量维度的张量进行扩展，最终生成5D张量，不符合模型要求的4D输入形状[1, height, width, 3]。

修正方案 #

移除手动的张量预处理和批量维度添加步骤，直接将canvas传入model.segment()，同时调整后续掩码与原图像的匹配逻辑：

import * as tf from '@tensorflow/tfjs';
import * as deeplab from '@tensorflow-models/deeplab';

export default async function rembg(imageUrl) {
  const model = await deeplab.load();
  const img = document.createElement('img');
  img.src = imageUrl;
  await img.decode();

  // 调整图像尺寸到模型兼容的输入大小
  const targetHeight = 513;
  const targetWidth = 513;

  const canvas = document.createElement('canvas');
  canvas.width = targetWidth;
  canvas.height = targetHeight;
  const ctx = canvas.getContext('2d');
  ctx.drawImage(img, 0, 0, targetWidth, targetHeight);

  // 直接传入canvas，由model.segment自动处理张量转换和预处理
  const segmentation = await model.segment(canvas);
  const backgroundMask = segmentation.segmentationMap;
  
  // 将掩码转换为与原图像匹配的形状，并转为float类型用于乘法
  const foregroundMask = backgroundMask.notEqual(0).toFloat();

  // 从canvas获取原图像张量，进行掩码运算
  const inputTensor = tf.browser.fromPixels(canvas).toFloat();
  const maskedImage = inputTensor.mul(foregroundMask);

  // 将结果转换为base64
  const pixels = await tf.browser.toPixels(maskedImage);
  const resultCanvas = document.createElement('canvas');
  resultCanvas.width = targetWidth;
  resultCanvas.height = targetHeight;
  const resultCtx = resultCanvas.getContext('2d');
  const imageData = new ImageData(pixels, targetWidth, targetHeight);
  resultCtx.putImageData(imageData, 0, 0);
  
  // 清理张量避免内存泄漏
  tf.dispose([inputTensor, foregroundMask, maskedImage, backgroundMask]);
  
  return resultCanvas.toDataURL('image/png');
}

关键调整说明 #

• 移除手动的张量归一化和expandDims(0)操作，让model.segment()内部处理输入格式
• 新增张量清理逻辑，避免浏览器内存泄漏
• 调整掩码处理步骤，确保与原图像张量形状和数据类型匹配

内容的提问来源于stack exchange，提问作者code now