coco标签处理

抽取coco格式训练集中的一部分作为子训练集，用快速测试算法

from pycocotools.coco import COCO
import random
import copy
import json# 原始的 COCO 标注文件路径
original_annotations_path = '/workspace/datasets/soccer6/yolo/annotations/train.json'# 新的 COCO 标注文件保存路径
new_annotations_path = './mini_train.json'# 读取原始的 COCO 标注文件
coco_original = COCO(original_annotations_path)# 获取所有图像的 IDs
all_image_ids = coco_original.getImgIds()# 随机选择 1000 张图像的 IDs
selected_image_ids = random.sample(all_image_ids, 1000)# 创建一个新的 COCO 数据结构
coco_new = copy.deepcopy(coco_original)
coco_new.dataset['images'] = [img for img in coco_original.dataset['images'] if img['id'] in selected_image_ids]
coco_new.dataset['annotations'] = [ann for ann in coco_original.dataset['annotations'] if ann['image_id'] in selected_image_ids]# 更新图像和标注的 ID
image_id_mapping = {old_id: new_id for new_id, old_id in enumerate(selected_image_ids, start=1)}
annotation_id_mapping = {ann['id']: i + 1 for i, ann in enumerate(coco_new.dataset['annotations'])}for img in coco_new.dataset['images']:img['id'] = image_id_mapping[img['id']]for ann in coco_new.dataset['annotations']:ann['image_id'] = image_id_mapping[ann['image_id']]ann['id'] = annotation_id_mapping[ann['id']]# 保存新的 COCO 标注文件
coco_new_path = new_annotations_path
with open(coco_new_path, 'w') as f:json.dump(coco_new.dataset, f)

测试新保存的json文件格式是否正确

from pycocotools.coco import COCO
import cv2
import os# COCO 数据集的路径（annotations 和 images 目录所在的路径）
coco_data_path = '/workspace/datasets/soccer6/yolo/'# COCO 标注文件路径
annotations_file_path = './mini_train.json'# 加载 COCO 数据集
coco = COCO(annotations_file_path)# 输出目录，用于保存带有标注的图像
output_directory = './output'# 确保输出目录存在
os.makedirs(output_directory, exist_ok=True)# 遍历所有图像
for image_id in coco.getImgIds():# 加载图像信息image_info = coco.loadImgs(image_id)[0]# 读取图像image_path = coco_data_path + '/images/' + image_info['file_name']image = cv2.imread(image_path)# 获取图像的标签信息ann_ids = coco.getAnnIds(imgIds=image_id)annotations = coco.loadAnns(ann_ids)# 在图像上绘制边界框和标签for annotation in annotations:bbox = annotation['bbox']category_id = annotation['category_id']category_name = coco.loadCats(category_id)[0]['name']if category_name == 'person':# 绘制边界框cv2.rectangle(image, (int(bbox[0]), int(bbox[1])),(int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])), (255, 0, 0), 2)# 在边界框上方绘制标签cv2.putText(image, category_name, (int(bbox[0]), int(bbox[1]) - 5),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)else:cv2.rectangle(image, (int(bbox[0]), int(bbox[1])),(int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])), (0, 0, 255), 2)# 在边界框上方绘制标签cv2.putText(image, category_name, (int(bbox[0]), int(bbox[1]) - 5),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)# 构建输出图像的路径output_path = os.path.join(output_directory, f'annotated_{image_info["file_name"]}')# 保存带有标注的图像cv2.imwrite(output_path, image)# 提示：如果你想使用 Matplotlib 显示图像，可以在循环内添加显示的代码

读取指定图片的标注

from pycocotools.coco import COCO
import cv2# COCO 数据集的路径（annotations 和 images 目录所在的路径）
coco_data_path = '/workspace/datasets/soccer6/yolo/'# 加载 COCO 数据集
coco = COCO('/workspace/datasets/soccer6/yolo/annotations/train.json')# 图像文件名
image_filename = 'groundball_373.jpg'# 获取图像信息
image_info = [img for img in coco.dataset['images'] if img['file_name'] == image_filename][0]# 读取图像
image_path = coco_data_path + '/images/' + image_info['file_name']
image = cv2.imread(image_path)# 获取图像的标签信息
ann_ids = coco.getAnnIds(imgIds=image_info['id'])
annotations = coco.loadAnns(ann_ids)# 打印标签信息
# 在图像上绘制边界框和标签
for annotation in annotations:bbox = annotation['bbox']category_id = annotation['category_id']category_name = coco.loadCats(category_id)[0]['name']# 绘制边界框cv2.rectangle(image, (int(bbox[0]), int(bbox[1])),(int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])), (255, 0, 0), 2)# 在边界框上方绘制标签cv2.putText(image, category_name, (int(bbox[0]), int(bbox[1]) - 5),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)# 保存带有标注的图像
output_path = 'annotated_image.jpg'
cv2.imwrite(output_path, image)