【yolo标签格式（txt）转coco格式(json)】

1.YOLO格式文件组织

这是要转化的YOLO格式文件组织，分别有images和labels文件，里面划分了train和val。

2.代码实现

将其转化为coco格式。

import os
import cv2
import json
import argparse
from tqdm import tqdm

COCO_DICT = ['images', 'annotations', 'categories']
IMAGES_DICT = ['file_name', 'height', 'width', 'id']
ANNOTATIONS_DICT = ['image_id', 'iscrowd', 'area', 'bbox', 'category_id', 'id']
CATEGORIES_DICT = ['id', 'name']

YOLO_CATEGORIES = [
    'Double hexagonal column', 'Flange nut', 'Hexagon nut', 'Hexagon pillar', 'Hexagon screw',
    'Hexagonal steel column', 'Horizontal bubble', 'Keybar', 'Plastic cushion pillar',
    'Rectangular nut', 'Round head screw', 'Spring washer', 'T-shaped screw'
]#这里替换成自己的类别

parser = argparse.ArgumentParser(description='Convert YOLO to COCO format')
parser.add_argument('--image_path', type=str, default='/home/yu/Yolov8/ultralytics-main/mydata0/images', help='Image folder containing train and val subfolders')#这里换成images的路径，images路径下包含train和val文件夹
parser.add_argument('--annotation_path', type=str, default='/home/yu/Yolov8/ultralytics-main/mydata0/labels', help='Annotation folder containing train and val subfolders')#这里换成标签的路径，同样包含标签的train和val文件夹
parser.add_argument('--save_dir', type=str, default='/home/yu/Yolov8/ultralytics-main/mydata0', help='Directory to save COCO JSON files')

args = parser.parse_args()


def save_json(data, path):
    print(f'Saving JSON: {path}')
    with open(path, 'w') as f:
        json.dump(data, f, indent=4)
    print('Successfully saved JSON:', path)


def load_image(image_path):
    img = cv2.imread(image_path)
    if img is None:
        print(f"Warning: Cannot load image {image_path}")
        return None, None
    return img.shape[0], img.shape[1]


def generate_categories_dict(category_list):
    return [{CATEGORIES_DICT[0]: i + 1, CATEGORIES_DICT[1]: name} for i, name in enumerate(category_list)]


def generate_images_dict(image_list, image_dir, start_image_id=0):
    images_dict = []
    for i, image_name in enumerate(tqdm(image_list, desc="Processing images")):
        image_path = os.path.join(image_dir, image_name)
        height, width = load_image(image_path)
        if height is None or width is None:
            continue  # Skip if image failed to load
        images_dict.append({
            'file_name': image_name,
            'height': height,
            'width': width,
            'id': start_image_id + i
        })
    return images_dict


def convert_yolo_to_coco(image_dir, label_dir, start_image_id=0, start_annotation_id=0):
    categories_dict = generate_categories_dict(YOLO_CATEGORIES)
    image_list = os.listdir(image_dir)
    images_dict = generate_images_dict(image_list, image_dir, start_image_id)

    annotations_dict = []
    annotation_id = start_annotation_id

    print("Processing annotations...")
    for image_data in images_dict:
        image_id = image_data['id']
        image_name = image_data['file_name']
        width, height = image_data['width'], image_data['height']

        label_file = os.path.join(label_dir, image_name.replace('.jpg', '.txt'))
        if not os.path.exists(label_file):
            continue  # Skip if no label file

        with open(label_file, 'r') as f:
            lines = f.readlines()

        for line in lines:
            parts = line.strip().split()
            if len(parts) != 5:
                continue  # Skip invalid annotation

            category_id = int(parts[0]) + 1  # YOLO starts from 0, COCO from 1
            x, y, w, h = map(float, parts[1:])

            x_min = (x - w / 2) * width
            y_min = (y - h / 2) * height
            bbox_width = w * width
            bbox_height = h * height
            area = bbox_width * bbox_height

            annotations_dict.append({
                'image_id': image_id,
                'iscrowd': 0,
                'area': area,
                'bbox': [x_min, y_min, bbox_width, bbox_height],
                'category_id': category_id,
                'id': annotation_id
            })
            annotation_id += 1

    return {
        'images': images_dict,
        'annotations': annotations_dict,
        'categories': categories_dict
    }


if __name__ == '__main__':
    for subset in ['train', 'val']:
        image_dir = os.path.join(args.image_path, subset)
        label_dir = os.path.join(args.annotation_path, subset)
        save_path = os.path.join(args.save_dir, f'instances_{subset}2017.json')

        if not os.path.exists(image_dir) or not os.path.exists(label_dir):
            print(f"Skipping {subset} as folders do not exist.")
            continue

        print(f"Processing {subset} dataset...")
        coco_data = convert_yolo_to_coco(image_dir, label_dir)
        save_json(coco_data, save_path)

执行完代码后分别有train和val的json文件,如下。

将其组织成coco数据集格式，即可。

3.COCO格式文件组织