coco标签处理

1. 抽取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)

2. 测试新保存的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 显示图像，可以在循环内添加显示的代码

3. 读取指定图片的标注

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)