用python写一个脚本将voc格式的数据集转换成coco格式

以下是一个将voc格式数据集转换为coco格式的Python脚本示例：

import os
import xml.etree.ElementTree as ET
import json

# Set the paths for the input and output directories
input_dir = 'path/to/voc/dataset'
output_dir = 'path/to/coco/dataset'

# Create the output directory if it doesn't exist
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# Define the categories for the dataset
categories = [
    {
        'id': 1,
        'name': 'cat',
        'supercategory': ''
    },
    {
        'id': 2,
        'name': 'dog',
        'supercategory': ''
    }
]

# Define the dictionary for storing the images and annotations
coco_data = {
    'images': [],
    'annotations': [],
    'categories': categories
}

# Define a function for adding an image to the dataset
def add_image(file_name, image_id):
    image = {
        'id': image_id,
        'file_name': file_name,
        'width': 0, # Set the width and height to 0 for now
        'height': 0
    }
    coco_data['images'].append(image)

# Define a function for adding an annotation to the dataset
def add_annotation(image_id, bbox, category_id, annotation_id):
    annotation = {
        'id': annotation_id,
        'image_id': image_id,
        'category_id': category_id,
        'bbox': bbox,
        'area': bbox[2] * bbox[3],
        'iscrowd': 0
    }
    coco_data['annotations'].append(annotation)

# Loop over the VOC dataset and convert each annotation to COCO format
image_id = 0
annotation_id = 0
for file_name in os.listdir(input_dir):
    if file_name.endswith('.xml'):
        # Parse the XML file
        xml_file = os.path.join(input_dir, file_name)
        tree = ET.parse(xml_file)
        root = tree.getroot()

        # Get the image dimensions
        size = root.find('size')
        width = int(size.find('width').text)
        height = int(size.find('height').text)

        # Add the image to the dataset
        add_image(file_name[:-4] + '.jpg', image_id)
        coco_data['images'][image_id]['width'] = width
        coco_data['images'][image_id]['height'] = height

        # Add the annotations to the dataset
        for obj in root.findall('object'):
            category = obj.find('name').text
            category_id = 1 if category == 'cat' else 2 # Set the category ID based on the object name
            bbox = obj.find('bndbox')
            xmin = int(bbox.find('xmin').text)
            ymin = int(bbox.find('ymin').text)
            xmax = int(bbox.find('xmax').text)
            ymax = int(bbox.find('ymax').text)
            bbox = [xmin, ymin, xmax - xmin, ymax - ymin]
            add_annotation(image_id, bbox, category_id, annotation_id)
            annotation_id += 1

        image_id += 1

# Write the dataset to a JSON file
output_file = os.path.join(output_dir, 'coco_dataset.json')
with open(output_file, 'w') as f:
    json.dump(coco_data, f)

该脚本将VOC格式的数据集转换为COCO格式，并将结果保存在JSON文件中。在此示例中，我们将猫和狗作为类别。如果您的数据集包含其他类别，请相应地更新类别列表。