HICO-DET数据集实战：用Python解析anno_bbox.mat，快速提取人-物交互标注信息

HICO-DET数据集实战：用Python解析anno_bbox.mat，快速提取人-物交互标注信息

在计算机视觉领域，人-物交互（Human-Object Interaction, HOI）检测是一个极具挑战性的任务。HICO-DET作为该领域的标杆数据集，其标注文件anno_bbox.mat采用MATLAB格式存储，这对习惯使用Python的开发者来说是个不小的障碍。本文将手把手教你如何用Python高效解析这个文件，提取关键信息并转换为深度学习框架友好的格式。

1. 环境准备与数据加载

首先确保已安装必要的Python库：

pip install scipy numpy pandas

anno_bbox.mat文件采用MATLAB v7.3格式，需要使用h5py库而非传统的scipy.io来加载：

import h5py def load_mat(file_path): """加载MATLAB v7.3格式文件""" with h5py.File(file_path, 'r') as f: data = {k: [f[v[0]][()] for v in f[k]] for k in f.keys()} return data

注意：HICO-DET的标注文件可能因版本不同而采用不同格式。如果遇到加载错误，可尝试先用scipy.io.loadmat()，失败后再切换至h5py方案。

2. 理解数据结构与关键字段

加载后的数据结构包含三个主要部分：

其中hoi字段的结构最为复杂，每个样本包含：

• id: 动作类别索引
• bboxhuman: 人物边界框数组[N,4]
• bboxobject: 物体边界框数组[M,4]
• connection: 交互对数组[K,2]，表示人物与物体的配对索引
• invis: 可见性标记

3. 标注信息提取实战

3.1 基础信息提取

首先提取图像基本信息：

def extract_basic_info(mat_data, split='train'): bbox_data = mat_data[f'bbox_{split}'] return { 'filenames': [str(f[0]) for f in bbox_data['filename'][0]], 'sizes': [tuple(map(int, s)) for s in bbox_data['size'][0]], 'hois': bbox_data['hoi'][0] }

3.2 交互关系解析

核心的HOI信息提取函数：

def parse_hoi_annotations(hoi_data): annotations = [] for i in range(len(hoi_data['id'][0])): hoi_id = int(hoi_data['id'][0][i][0]) human_boxes = hoi_data['bboxhuman'][0][i].reshape(-1,4) if hoi_data['bboxhuman'][0][i].size else [] obj_boxes = hoi_data['bboxobject'][0][i].reshape(-1,4) if hoi_data['bboxobject'][0][i].size else [] connections = hoi_data['connection'][0][i].reshape(-1,2) if hoi_data['connection'][0][i].size else [] annotations.append({ 'hoi_id': hoi_id, 'human_boxes': human_boxes.tolist(), 'object_boxes': obj_boxes.tolist(), 'connections': connections.tolist() }) return annotations

3.3 动作类别映射

list_action包含600个HOI类别的详细信息：

def get_action_mapping(mat_data): return [{ 'id': i+1, 'verb': str(mat_data['list_action']['vname'][0][i][0]), 'object': str(mat_data['list_action']['nname'][0][i][0]) } for i in range(len(mat_data['list_action']['nname'][0]))]

4. 转换为COCO格式

为方便训练，我们将数据转换为COCO风格的JSON格式：

def convert_to_coco(mat_data, split='train'): # 初始化COCO结构 coco_data = { 'images': [], 'annotations': [], 'categories': get_action_mapping(mat_data) } basic_info = extract_basic_info(mat_data, split) for img_idx, (filename, size) in enumerate(zip(basic_info['filenames'], basic_info['sizes'])): # 添加图像信息 coco_data['images'].append({ 'id': img_idx, 'file_name': filename, 'width': size[0], 'height': size[1] }) # 处理该图像的所有HOI标注 hoi_anns = parse_hoi_annotations(basic_info['hois'][img_idx]) for hoi in hoi_anns: for conn in hoi['connections']: human_idx, obj_idx = conn human_box = hoi['human_boxes'][human_idx-1] # MATLAB索引从1开始 obj_box = hoi['object_boxes'][obj_idx-1] coco_data['annotations'].append({ 'image_id': img_idx, 'hoi_id': hoi['hoi_id'], 'human_bbox': human_box, 'object_bbox': obj_box, 'category_id': hoi['hoi_id'] }) return coco_data

5. 实用技巧与性能优化

处理大型MAT文件时，可采用以下优化策略：

1. 分块处理：对于特别大的文件，可以分块加载和处理数据

def process_in_chunks(mat_path, chunk_size=1000): with h5py.File(mat_path, 'r') as f: total = len(f['bbox_train/filename']) for i in range(0, total, chunk_size): chunk = {k: f[k][i:i+chunk_size] for k in f.keys()} yield chunk

2. 并行处理：利用多进程加速

from multiprocessing import Pool def parallel_parse(args): img_idx, hoi_data = args return parse_hoi_annotations(hoi_data) with Pool(4) as p: # 使用4个进程 results = p.map(parallel_parse, enumerate(basic_info['hois']))

3. 缓存机制：将解析结果保存为JSON，避免重复解析

import json from pathlib import Path def load_or_parse(mat_path, cache_dir='cache'): cache_path = Path(cache_dir) / (Path(mat_path).stem + '.json') if cache_path.exists(): return json.loads(cache_path.read_text()) data = load_mat(mat_path) coco_data = convert_to_coco(data) cache_path.parent.mkdir(exist_ok=True) cache_path.write_text(json.dumps(coco_data)) return coco_data

6. 可视化验证

为确保解析正确，建议可视化部分结果：

import cv2 import matplotlib.pyplot as plt def visualize_hoi(image_path, human_box, obj_box, action_name): img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB) # 绘制边界框 cv2.rectangle(img, (human_box[0], human_box[1]), (human_box[2], human_box[3]), (255,0,0), 2) cv2.rectangle(img, (obj_box[0], obj_box[1]), (obj_box[2], obj_box[3]), (0,255,0), 2) # 添加动作标签 plt.imshow(img) plt.title(action_name) plt.axis('off') plt.show()

7. 与深度学习框架集成

最后，我们可以创建PyTorch Dataset类方便训练使用：

from torch.utils.data import Dataset class HOIDataset(Dataset): def __init__(self, coco_data, image_dir, transform=None): self.data = coco_data self.image_dir = Path(image_dir) self.transform = transform def __len__(self): return len(self.data['images']) def __getitem__(self, idx): img_info = self.data['images'][idx] img_path = self.image_dir / img_info['file_name'] image = cv2.imread(str(img_path)) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 获取该图像的所有标注 anns = [a for a in self.data['annotations'] if a['image_id'] == idx] if self.transform: image = self.transform(image) return image, anns