如何从TensorFlowJS目标检测模型输出中提取bounding box、类别与得分？

你调用的SSD MobileNet V2 TF.js版本返回的两个张量分别对应检测得分和候选框，需要额外做非极大值抑制（NMS）过滤，就能拿到你要的边界框、类别标签、得分三个参数，修改代码如下：

function getMethods(o) {
    return Object.getOwnPropertyNames(Object.getPrototypeOf(o))
        .filter(m => 'function' === typeof o[m]);
}

// COCO数据集类别标签映射，可根据需要补全全量90类标签
const COCO_CLASSES = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow'
    // 剩余类别可自行搜索COCO类别列表补全
]

const model = await tf.loadGraphModel("https://tfhub.dev/tensorflow/tfjs-model/ssd_mobilenet_v2/1/default/1", {fromTFHub: true});
console.log("Get methods");
console.log(getMethods(model));

const cat = document.getElementById('cat');
let tensor = tf.browser.fromPixels(cat);
tensor = tf.cast(tensor, 'int32');
let y_pred = await model.executeAsync({'image_tensor': tensor.expandDims(0)});

// 解构返回结果：第一个是得分张量，第二个是边界框张量
const [scoresTensor, boxesTensor] = y_pred;

// 去掉张量多余的空维度
const scores = tf.squeeze(scoresTensor);
const boxes = tf.squeeze(boxesTensor);

// 执行非极大值抑制，过滤重复、低置信度的候选框
const nmsResults = await tf.image.nonMaxSuppressionAsync(
    boxes, // 所有候选框
    tf.max(scores, 1), // 每个候选框的最高置信度
    20, // 最多返回的检测框数量
    0.5, // IOU过滤阈值
    0.5 // 最低置信度阈值
);

// 将张量转为普通JS数组方便后续处理
const boxesData = await boxes.array();
const scoresData = await scores.array();
const validIndices = await nmsResults.array();

// 遍历有效索引组装最终检测结果
const finalDetections = validIndices.map(idx => {
    const [ymin, xmin, ymax, xmax] = boxesData[idx];
    // 将归一化坐标转为图片实际像素坐标
    const imgWidth = cat.width;
    const imgHeight = cat.height;
    const bbox = {
        x: Math.round(xmin * imgWidth),
        y: Math.round(ymin * imgHeight),
        width: Math.round((xmax - xmin) * imgWidth),
        height: Math.round((ymax - ymin) * imgHeight)
    };
    const maxScore = Math.max(...scoresData[idx]);
    const classId = scoresData[idx].indexOf(maxScore);
    const className = COCO_CLASSES[classId] || 'unknown';
    return {
        bbox,
        className,
        score: maxScore
    }
});

console.log('最终检测结果：', finalDetections);

// 使用完成后释放张量内存，避免页面内存泄漏
tf.dispose([tensor, y_pred, scoresTensor, boxesTensor, scores, boxes, nmsResults]);

补充说明 #

• 返回的第一个张量形状[1, 1917, 90]，其中90对应COCO数据集的90个类别，每个候选框对应90个类别的置信度
• 返回的第二个张量形状[1, 1917, 1, 4]，最后一维的4个值是归一化的边界框坐标[ymin, xmin, ymax, xmax]，范围在0~1之间，需要乘图片实际宽高才能得到对应像素坐标
• 非极大值抑制是必须步骤，用来过滤模型输出的大量重复候选框，只保留置信度最高的有效结果

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