如何在Detectron2中统计语义分割各类别总像素数量

Detectron2语义分割输出的sem_seg字段是形状为[C, H, W]的Tensor（C为类别数，H、W是图像高宽），每个通道对应对应类别的预测概率。要统计每个通用类别的总像素数，只需先确定每个像素的预测类别，再统计各类别像素数量，具体步骤如下：

实现代码 #

import torch

# 从语义分割输出中获取sem_seg张量
sem_seg = output['sem_seg']

# 对每个像素取概率最大的类别索引，得到[H, W]的类别预测图
pred_classes = torch.argmax(sem_seg, dim=0)

# 展平张量后统计每个类别的像素数量
flat_pred = pred_classes.flatten()
class_pixel_counts = torch.bincount(flat_pred)

代码说明 #

• torch.argmax(sem_seg, dim=0)：沿着通道维度（dim=0）取最大值的索引，每个像素位置会得到其所属的类别ID，最终生成和原图尺寸一致的类别预测矩阵。
• torch.bincount(flat_pred)：对展平后的类别ID数组进行计数，返回的class_pixel_counts是一维Tensor，其中class_pixel_counts[i]就代表类别ID为i的总像素数。

扩展说明 #

如果需要只统计特定类别（比如仅统计“人”“车辆”类），可以先筛选出目标类别对应的像素再计数：

# 假设目标类别ID是1（人）和2（车辆）
target_classes = [1, 2]
for cls_id in target_classes:
    pixel_count = torch.sum(pred_classes == cls_id).item()
    print(f"类别ID {cls_id} 的总像素数：{pixel_count}")

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