VOC数据集转COCO数据集工具

当我们处理目标检测任务时，很多时候我们会遇到必须将数据集从 VOC 格式转换为 COCO 格式的情况。VOC 格式和 COCO 格式是两种广泛使用的目标检测数据集格式。VOC 格式使用 XML 文件来存储每个图像的标注信息，而 COCO 格式使用 JSON 文件。这种格式转换通常是为了适应不同的深度学习框架或工具。

为了简化这个过程，我给大家分享一个 Python 脚本，可以将 VOC 格式的数据集转换为 COCO 格式，并且还支持自动复制图像到指定目录。

首先，我们需要准备以下内容：

• VOC 格式的标注文件夹（包含 XML 文件）
• 目标转换后的 COCO 格式 JSON 文件所在文件夹
• 类别列表（以字符串形式表示，比如 ['cat', 'dog', 'person']）
• 图像文件夹，存放对应的图像文件

接下来，我们创建一个 VOC2COCOConverter 对象，并指定上面提到的各个参数。我们还可以选择设置 proportions 参数，以控制最后生成的 COCO 格式数据集在训练集、验证集和测试集之间的比例。默认情况下，这个参数设置为 [80, 10, 10]，即将数据集划分为 80%的训练集，10%的验证集和 10%的测试集。当然如果你需要也可以设置为 [100] 或 [80, 20]。

我们还可以设置 copy_images 参数来决定是否复制图像文件。如果将其设置为 True，则脚本会自动将图像复制到与生成的 COCO 格式 JSON 文件同名的文件夹中。这个功能对于数据集管理很有用，可以在用其他框架操作数据集、数据探索或模型调试时方便地使用数据集。

下面是示例代码的部分内容：

import os
import glob
import json
import shutil
import xml.etree.ElementTree as ET
from collections import defaultdict, Counter
from tqdm import tqdm

START_BOUNDING_BOX_ID = 1

class VOC2COCOConverter:
    def __init__(self, xml_dir, json_dir, classes, img_dir, proportions=[8, 1, 1], copy_images=False, min_samples_per_class=20):
        self.xml_dir = xml_dir
        self.json_dir = json_dir
        self.img_dir = img_dir
        self.classes = classes
        self.proportions = proportions
        self.copy_images = copy_images
        self.min_samples_per_class = min_samples_per_class

        self.pre_define_categories = {}
        for i, cls in enumerate(self.classes):
            self.pre_define_categories[cls] = i + 1
            
    def convert(self):
        xml_files_by_class = self._get_sorted_xml_files_by_class()
        dataset_size = len(self.proportions)
        xml_files_by_dataset = [defaultdict(list) for _ in range(dataset_size)]
        xml_files_count_by_dataset = [0] * dataset_size
        
        for cls, xml_files in xml_files_by_class.items():
            total_files = len(xml_files)
            datasets_limits = [int(total_files * p / sum(self.proportions)) for p in self.proportions]
            datasets_limits[-1] = total_files - sum(datasets_limits[:-1]) # adjust to make sure the sums are correct due to integer division

            start = 0
            for i, limit in enumerate(datasets_limits):
                xml_files_by_dataset[i][cls] = xml_files[start:start + limit]
                xml_files_count_by_dataset[i] += limit
                start += limit

        for idx, xml_files_dict in enumerate(xml_files_by_dataset):
            dataset_dir = ''
            if self.copy_images:
                dataset_dir = os.path.join(self.json_dir, f'dataset_{idx + 1}')
                os.makedirs(dataset_dir, exist_ok=True)
                
            json_file_name = f'dataset_{idx + 1}.json'
            xml_files = sum(xml_files_dict.values(), [])
            self._convert_annotation(tqdm(xml_files), os.path.join(self.json_dir, json_file_name))
            if dataset_dir:
                self._copy_images(tqdm(xml_files), dataset_dir)

            print(f"\n在数据集{idx+1}中，各个类型的样本数量分别为：")
            for cls, files in xml_files_dict.items():
                print(f"类型 {cls} 的样本数量是: {len(files)}")

        print("\n各个数据集中相同类型样本的数量比值是：")
        for cls in self.classes:
            print("\n类型 {}:".format(cls))
            for i in range(len(self.proportions) - 1):
                if len(xml_files_by_dataset[i + 1].get(cls, [])) != 0 :
                    print("数据集 {} 和 数据集 {} 的样本数量比是: {}".format(
                        i + 1, 
                        i + 2, 
                        len(xml_files_by_dataset[i].get(cls, [])) / len(xml_files_by_dataset[i + 1].get(cls, []))
                    ))

    def _get_sorted_xml_files_by_class(self):
        xml_files_by_class = defaultdict(list)
        for xml_file in glob.glob(os.path.join(self.xml_dir, "*.xml")):
            tree = ET.parse(xml_file)
            root = tree.getroot()
            for obj in root.findall('object'):
                class_name = obj.find('name').text
                if class_name in self.classes:
                    xml_files_by_class[class_name].append(xml_file)


        # Filter classes
        if self.min_samples_per_class is not None:
            xml_files_by_class = {
                cls: files 
                for cls, files in xml_files_by_class.items()
                if len(files) > self.min_samples_per_class
            }

        xml_files_by_class = dict(
            sorted(xml_files_by_class.items(), key=lambda item: len(item[1]), reverse=True))

        return xml_files_by_class

    def _copy_images(self, xml_files, dataset_dir):
        for xml_file in xml_files:
            img_file = os.path.join(self.img_dir, os.path.basename(xml_file).replace('.xml', '.jpg'))
            if os.path.exists(img_file):
                shutil.copy(img_file, dataset_dir)

    def _get_files_by_majority_class(self):
        xml_files_by_class = defaultdict(list)
        for xml_file in glob.glob(os.path.join(self.xml_dir, "*.xml")):
            tree = ET.parse(xml_file)
            root = tree.getroot()
            class_counts = defaultdict(int)
            for obj in root.findall('object'):
                class_name = obj.find('name').text
                if class_name in self.classes:
                    class_counts[class_name] += 1
            majority_class = max(class_counts, key=class_counts.get)
            xml_files_by_class[majority_class].append(xml_file)

        return dict(sorted(xml_files_by_class.items(), key=lambda item: len(item[1]), reverse=True))

    def _convert_annotation(self, xml_list, json_file):
        json_dict = {"info":['none'], "license":['none'], "images": [], "annotations": [], "categories": []}
        categories = self.pre_define_categories.copy()
        bnd_id = START_BOUNDING_BOX_ID
        all_categories = {}

        for index, line in enumerate(xml_list):
            xml_f = line
            tree = ET.parse(xml_f)
            root = tree.getroot()

            filename = os.path.basename(xml_f)[:-4] + ".jpg"

            image_id = int(filename.split('.')[0][-9:])

            size = self._get_and_check(root, 'size', 1)
            width = int(self._get_and_check(size, 'width', 1).text)
            height = int(self._get_and_check(size, 'height', 1).text)
            image = {'file_name': filename, 'height': height, 'width': width, 'id':image_id}
            json_dict['images'].append(image)

            for obj in self._get(root, 'object'):
                category = self._get_and_check(obj, 'name', 1).text
                if category in all_categories:
                    all_categories[category] += 1
                else:
                    all_categories[category] = 1

                if category not in categories:
                    new_id = len(categories) + 1
                    print(filename)
                    print("[warning] 类别 '{}' 不在 'pre_define_categories'({})中，将自动创建新的id: {}".format(category, self.pre_define_categories, new_id))
                    categories[category] = new_id

                category_id = categories[category]
                bndbox = self._get_and_check(obj, 'bndbox', 1)
                xmin = int(float(self._get_and_check(bndbox, 'xmin', 1).text))
                ymin = int(float(self._get_and_check(bndbox, 'ymin', 1).text))
                xmax = int(float(self._get_and_check(bndbox, 'xmax', 1).text))
                ymax = int(float(self._get_and_check(bndbox, 'ymax', 1).text))
                o_width = abs(xmax - xmin)
                o_height = abs(ymax - ymin)

                ann = {'area': o_width*o_height, 'iscrowd': 0, 'image_id': image_id, 'bbox':[xmin, ymin, o_width, o_height],
                       'category_id': category_id, 'id': bnd_id, 'ignore': 0, 'segmentation': []}
                json_dict['annotations'].append(ann)
                bnd_id = bnd_id + 1

        for cate, cid in categories.items():
            cat = {'supercategory': 'none', 'id': cid, 'name': cate}
            json_dict['categories'].append(cat)
        json_fp = open(json_file, 'w')
        json_str = json.dumps(json_dict)
        json_fp.write(json_str)
        json_fp.close()
        print("------------已完成创建 {}--------------".format(json_file))
        print("找到 {} 类别: {} -->>> 你的预定类别 {}: {}".format(len(all_categories), all_categories.keys(), len(self.pre_define_categories), self.pre_define_categories.keys()))
        print("类别: id --> {}".format(categories))
        
    def _get(self, root, name):
        return root.findall(name)

    def _get_and_check(self, root, name, length):
        vars = root.findall(name)
        if len(vars) == 0:
            raise NotImplementedError('Can not find %s in %s.'%(name, root.tag))
        if length > 0 and len(vars) != length:
            raise NotImplementedError('The size of %s is supposed to be %d, but is %d.'%(name, length, len(vars)))
        if length == 1:
            vars = vars[0]
        return vars

if __name__ == '__main__':
    # xml标注文件夹   
    xml_dir = 'path/to/xml/directory'
    # JSON文件所在文件夹
    json_dir = 'path/to/json/directory'  
    # 类别列表，以字符串形式表示
    classes = ['cat', 'dog', 'person']
    # 图片所在文件夹
    img_dir = 'path/to/image/directory'
    # 类别在数据集中的比例
    proportions = [80, 10, 10]
    # 创建VOC2COCOConverter对象并进行转换
    converter = VOC2COCOConverter(xml_dir, json_dir, classes, img_dir, proportions, copy_images=True)
    converter.convert()

上面的代码只是将voc格式转换成coco格式，并不会指定具体哪个数据集的作用，因此在进行转换后需要手动给每个文件夹与注释文件命名，其中test数据集不是必要的，可以根据你的需要进行调整。

这是COCO数据集的基本目录结构：

|-- annotations
|   |-- instances_train.json
|   |-- instances_val.json
|   |-- instances_test.json
|-- train
|   |-- image1.jpg
|   |-- image2.jpg
|   |-- ...
|-- val
|   |-- image1.jpg
|   |-- image2.jpg
|   |-- ...
|-- test
|   |-- image1.jpg
|   |-- image2.jpg
|   |-- ...

annotations文件夹：存放标注文件，如instances_train.json、instances_val.json等。

train文件夹：存放训练集的图像文件。

val文件夹：存放验证集的图像文件。

test文件夹：存放测试集的图像文件。

以上是这个 VOC 到 COCO 格式转换脚本的简单介绍。你可以根据自己的实际需求进行适当修改和优化。这个脚本可以帮助你高效地转换数据集格式，并且支持自动复制图像文件，方便管理和使用数据集。

注意：本文及代码完全由 GPT 4 自动生成，经测试代码可以正常使用。