COCO目标检测数据集的读取方法与Python工具脚本
最新回答
栖止你掌
2021-07-14 18:01:18
COCO数据集格式介绍
COCO (Common Objects in COntext) 是一个大型图像数据集，提供目标检测、分割、看图说话等任务的标注。其标注文件采用JSON格式，主要包含以下结构：
{ "info": info, "images": [image], "annotations": [annotation], "licenses": [license], "categories": [category]}关键字段说明图像信息 (image)
{ "id": int, "width": int, "height": int, "file_name": str, "license": int, "flickr_url": str, "coco_url": str, "date_captured": datetime}标注信息 (annotation)
{ "id": int, "image_id": int, "category_id": int, "segmentation": RLE or [polygon], "area": float, "bbox": [x,y,width,height], "iscrowd": 0 or 1}类别信息 (category)
{ "id": int, "name": str, "supercategory": str}
Python读取脚本
基本读取方法import jsondef load_coco_json(ann_path): """加载COCO格式的JSON标注文件""" with open(ann_path) as fp: root = json.load(fp) return root# 示例用法root = load_coco_json('data/coco/annotations/instances_val2014.json')获取关键信息def print_coco_info(root): """打印COCO标注文件的基本信息""" print('Dataset Info:') print(root['info']) print('nCategories:') for category in root['categories']: print(f"ID: {category['id']}, Name: {category['name']}") print(f'nTotal Images: {len(root["images"])}') print(f'Total Annotations: {len(root["annotations"])}')# 示例用法print_coco_info(root)构建图像-标注字典def build_img_ann_dict(ann_path): """构建图像ID到标注信息的字典 返回格式: { img_id: { 'file_name': str, 'annotations': [ [x, y, width, height, category_id], ... ] } } """ with open(ann_path) as fp: root = json.load(fp) img_dict = {} # 首先构建图像ID到文件名的映射 for img_info in root['images']: img_dict[img_info['id']] = { 'file_name': img_info['file_name'], 'annotations': [] } # 填充标注信息 for ann in root['annotations']: img_id = ann['image_id'] bbox = ann['bbox'] category_id = ann['category_id'] img_dict[img_id]['annotations'].append(bbox + [category_id]) return img_dict# 示例用法img_dict = build_img_ann_dict('data/coco/annotations/instances_val2014.json')类别ID到名称映射def build_category_map(root): """构建类别ID到名称的映射字典""" return {category['id']: category['name'] for category in root['categories']}# 示例用法category_map = build_category_map(root)print(category_map[1]) # 输出: 'person'完整示例脚本
import jsonimport osfrom PIL import Imageclass COCODataset: def __init__(self, ann_path, img_dir): """初始化COCO数据集读取器 Args: ann_path: 标注文件路径 img_dir: 图像目录路径 """ self.ann_path = ann_path self.img_dir = img_dir self.root = self._load_json() self.img_dict = self._build_img_dict() self.category_map = self._build_category_map() def _load_json(self): """加载JSON标注文件""" with open(self.ann_path) as fp: return json.load(fp) def _build_img_dict(self): """构建图像字典""" img_dict = {} # 构建图像ID到文件名的映射 for img_info in self.root['images']: img_dict[img_info['id']] = { 'file_name': img_info['file_name'], 'annotations': [] } # 填充标注信息 for ann in self.root['annotations']: img_id = ann['image_id'] bbox = ann['bbox'] category_id = ann['category_id'] img_dict[img_id]['annotations'].append({ 'bbox': bbox, 'category_id': category_id }) return img_dict def _build_category_map(self): """构建类别映射字典""" return {category['id']: category['name'] for category in self.root['categories']} def get_image_info(self, img_id): """获取指定图像的信息""" return self.img_dict.get(img_id) def visualize_annotations(self, img_id, output_path='output.jpg'): """可视化指定图像的标注 Args: img_id: 图像ID output_path: 输出图像路径 """ img_info = self.img_dict.get(img_id) if not img_info: raise ValueError(f"Image with ID {img_id} not found") # 加载图像 img_path = os.path.join(self.img_dir, img_info['file_name']) img = Image.open(img_path) # 绘制每个标注框 for ann in img_info['annotations']: bbox = ann['bbox'] category_id = ann['category_id'] x, y, w, h = bbox x1, y1, x2, y2 = x, y, x + w, y + h # 这里可以添加实际的绘制代码 # 示例中省略了实际的绘制实现 print(f"Drawing bbox: ({x1},{y1})-({x2},{y2}), Category: {self.category_map[category_id]}") # 保存可视化结果 img.save(output_path) print(f"Visualization saved to {output_path}")# 示例用法if __name__ == '__main__': # 初始化数据集读取器 dataset = COCODataset( ann_path='data/coco/annotations/instances_val2014.json', img_dir='data/coco/val2014' ) # 获取第一个图像的信息 img_ids = list(dataset.img_dict.keys()) first_img_info = dataset.get_image_info(img_ids[0]) print(f"First image: {first_img_info['file_name']}") print(f"Number of annotations: {len(first_img_info['annotations'])}") # 可视化标注 dataset.visualize_annotations(img_ids[0], 'work_dirs/visualization.jpg')注意事项
类别ID不连续：COCO数据集中的类别ID不是从0开始连续的，最大ID为90但实际只有80个类别。在实际应用中需要建立0-79到实际category_id的映射。
坐标系统：COCO的坐标系统原点在图像左上角，x向右增长，y向下增长。
iscrowd字段：当iscrowd=1时，表示该标注对应一组对象（如人群），此时segmentation字段使用RLE格式编码。
性能优化：对于大型数据集，可以考虑使用更高效的数据结构或缓存机制来存储和访问标注信息。
可视化实现：完整示例中的可视化部分省略了实际的绘制代码，实际应用中可以使用OpenCV或PIL等库实现标注框的绘制。
扩展功能
数据分割：可以根据需要实现训练集、验证集和测试集的分割。
数据增强：在读取数据时可以集成数据增强功能。
多任务支持：可以扩展支持实例分割、关键点检测等其他COCO任务。
懒加载：对于大型数据集，可以实现按需加载图像和标注的懒加载机制。
热门标签
