如何在Python/Numpy中提取同类别ID的目标检测框索引与数据

嘿，我来帮你搞定这个问题！首先得明确：你给出的n×4数组是检测框的(x, y, width, height)坐标，但类别ID（class id）应该是和检测框一一对应的独立数组（毕竟示例里没包含class信息）。下面我分两种场景给你实现方案：

场景1：纯Python实现（小数据量友好） #

如果你的检测框和class ID都是普通列表，用字典分组是最直观的方式，键为class ID，值存储对应的（索引，检测框数据）元组：

# 示例检测框列表
detections = [
    [2433, 1629, 2656, 1910],
    [2275, 962, 2538, 1264],
    [2284, 1191, 2533, 1472],
    [2303, 1380, 2551, 1662],
    [2461, 1442, 2694, 1732],
    [2284, 1558, 2533, 1852],
    [2433, 1261, 2671, 1524]
]
# 对应每个检测框的class ID列表（长度和detections一致）
class_ids = [0, 1, 1, 0, 2, 1, 0]

# 初始化分组字典
class_groups = {}

# 遍历检测框和class ID，按class分组
for idx, (box, cls_id) in enumerate(zip(detections, class_ids)):
    if cls_id not in class_groups:
        class_groups[cls_id] = []
    class_groups[cls_id].append( (idx, box) )

# 打印结果
for cls_id, items in class_groups.items():
    print(f"=== Class ID {cls_id} ===")
    for idx, box in items:
        print(f"索引 {idx}: 检测框 {box}")

运行后会输出每个class ID对应的所有目标索引和检测框数据。

场景2：Numpy实现（大数据量高效） #

如果你的数据量很大（比如上万条检测框），用Numpy的向量操作会比纯Python循环快很多。步骤是：将数据转为Numpy数组，利用布尔索引筛选对应class的目标：

import numpy as np

# 转为Numpy数组
detections_np = np.array(detections)
class_ids_np = np.array(class_ids)

# 获取所有唯一的class ID
unique_classes = np.unique(class_ids_np)

# 分组存储结果
class_groups_np = {}

for cls_id in unique_classes:
    # 筛选该class对应的所有索引
    target_indices = np.where(class_ids_np == cls_id)[0]
    # 筛选对应的检测框
    target_boxes = detections_np[target_indices]
    # 存储索引和检测框（转为列表方便查看）
    class_groups_np[cls_id] = list(zip(target_indices, target_boxes.tolist()))

# 打印结果
for cls_id, items in class_groups_np.items():
    print(f"=== Class ID {cls_id} ===")
    for idx, box in items:
        print(f"索引 {idx}: 检测框 {box}")

额外情况：如果检测输出包含Class ID #

有些目标检测器的输出是n×5的数组（最后一列是class ID），这种情况只需要先拆分出检测框和class ID即可：

# 示例n×5的检测输出
full_detections_np = np.array([
    [2433, 1629, 2656, 1910, 0],
    [2275, 962, 2538, 1264, 1],
    [2284, 1191, 2533, 1472, 1],
    # ...更多数据
])

# 拆分检测框和class ID
boxes_np = full_detections_np[:, :4]  # 前4列是坐标
class_ids_np = full_detections_np[:, 4]  # 最后一列是class ID

# 后续处理和上面的Numpy方案完全一致

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