如何基于Open MMLAB的MMDetection筛选含person标签的图片？

方案一：推理后过滤结果（适用于单张/批量图片推理场景） #

COCO数据集标注中person的类别ID为1，对应MMDetection模型输出的类别索引为0（模型训练时会将原始类别ID映射为从0开始的索引）。可以在推理完成后，检查每张图片的检测结果中是否包含person类别目标，以此判断是否保留该图片。

具体步骤及代码示例：

• 加载预训练模型与推理工具
• 遍历目标图片执行推理
• 解析结果，判断是否存在置信度达标（可自定义阈值）的person检测框
• 对符合条件的图片进行留存处理

from mmdet.apis import init_detector, inference_detector
import os
import shutil

# 配置文件与预训练权重路径
config_file = 'configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
checkpoint_file = 'checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
device = 'cuda:0'

# 初始化模型
model = init_detector(config_file, checkpoint_file, device=device)

# 待处理目录与保存目录
img_dir = 'path/to/your/images'
save_dir = 'path/to/save/person_images'
os.makedirs(save_dir, exist_ok=True)

# 遍历图片筛选
for img_name in os.listdir(img_dir):
    img_path = os.path.join(img_dir, img_name)
    result = inference_detector(model, img_path)
    # 提取person类别检测结果，置信度阈值设为0.5
    person_boxes = result[0]
    if len(person_boxes) > 0 and (person_boxes[:, -1] >= 0.5).any():
        shutil.copy(img_path, os.path.join(save_dir, img_name))

方案二：数据集加载阶段筛选（适用于训练/验证数据集预处理场景） #

如果是要从COCO格式的数据集中提前筛选含person标签的样本，可通过两种方式实现：

方法1：修改MMDetection数据集配置 #

在数据集配置中添加自定义过滤函数，仅保留含person标注的样本：

def filter_person_samples(data_info):
    # 检查样本标注是否包含person类别（COCO标注中person的category_id为1）
    has_person = any(ann['category_id'] == 1 for ann in data_info['annotations'])
    return has_person

# 配置文件中应用过滤函数
data = dict(
    train=dict(
        type='CocoDataset',
        ann_file='data/coco/annotations/instances_train2017.json',
        img_prefix='data/coco/train2017/',
        pipeline=train_pipeline,
        filter_empty_gt=False,
        filter_fn=filter_person_samples
    ),
    # val、test数据集可同理配置
)

方法2：预处理COCO标注文件 #

直接修改原始COCO JSON标注文件，仅保留含person的图片与对应标注：

import json

# 加载原始标注文件
with open('data/coco/annotations/instances_train2017.json', 'r') as f:
    coco_data = json.load(f)

# 筛选含person的图片ID
person_category_id = 1
person_img_ids = set()
for ann in coco_data['annotations']:
    if ann['category_id'] == person_category_id:
        person_img_ids.add(ann['image_id'])

# 过滤图片、标注与类别
filtered_images = [img for img in coco_data['images'] if img['id'] in person_img_ids]
filtered_annotations = [ann for ann in coco_data['annotations'] if ann['image_id'] in person_img_ids]
filtered_categories = [cat for cat in coco_data['categories'] if cat['id'] == person_category_id]

# 保存新标注文件
filtered_coco_data = {
    'info': coco_data['info'],
    'licenses': coco_data['licenses'],
    'categories': filtered_categories,
    'images': filtered_images,
    'annotations': filtered_annotations
}

with open('data/coco/annotations/instances_train2017_person_only.json', 'w') as f:
    json.dump(filtered_coco_data, f)

之后在MMDetection配置文件中使用该新标注文件即可。

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