Cuando nos ocupamos de tareas de detección de objetos, muchas veces nos encontramos con situaciones en las que debemos convertir el conjunto de datos del formato VOC al formato COCO. El formato VOC y el formato COCO son dos formatos de conjuntos de datos de detección de objetivos ampliamente utilizados. El formato VOC utiliza archivos XML para almacenar información de anotaciones para cada imagen, mientras que el formato COCO utiliza archivos JSON. Esta conversión de formato generalmente se realiza para adaptarse a diferentes marcos o herramientas de aprendizaje profundo.
Para simplificar este proceso, compartiré con ustedes un script de Python que puede convertir conjuntos de datos en formato VOC al formato COCO y también admite la copia automática de imágenes a un directorio específico.
Primero, necesitamos preparar lo siguiente:
- Carpeta de anotaciones en formato VOC (contiene archivos XML)
- La carpeta donde se encuentra el archivo JSON en formato COCO convertido de destino
- lista de categorías (representadas como una cadena, por ejemplo
['cat', 'dog', 'person']) - Carpeta de imágenes, que almacena los archivos de imagen correspondientes.
A continuación, creamos un VOC2COCOConverterobjeto y especificamos los diversos parámetros mencionados anteriormente. También podemos optar por establecer proportionsparámetros para controlar la proporción del conjunto de datos en formato COCO final generado entre el conjunto de entrenamiento, el conjunto de validación y el conjunto de prueba. De forma predeterminada, este parámetro está configurado [80, 10, 10]para dividir el conjunto de datos en un 80% de conjunto de entrenamiento, un 10% de conjunto de validación y un 10% de conjunto de prueba. [100]Por supuesto, también puedes configurarlo en o si lo necesitas [80, 20].
También podemos establecer copy_imagesparámetros para decidir si copiar archivos de imagen. Si está configurado en True, el script copia automáticamente la imagen en una carpeta con el mismo nombre que el archivo JSON en formato COCO generado. Esta característica es útil para la gestión de conjuntos de datos y se puede utilizar cómodamente cuando se utilizan otros marcos para manipular conjuntos de datos, explorar datos o depurar modelos.
Aquí está parte del código de muestra:
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()
El código anterior solo convierte el formato voc al formato coco y no especifica qué conjunto de datos se utiliza. Por lo tanto, después de la conversión, debe nombrar manualmente cada carpeta y archivo de anotación. El conjunto de datos de prueba no es necesario y puede basarse en Ajustar a tus necesidades.
Esta es la estructura de directorio básica del conjunto de datos 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
| |-- ...
Carpeta de anotaciones: almacena archivos de anotaciones, como instancias_train.json, instancias_val.json, etc.
Carpeta de tren: almacena los archivos de imagen del conjunto de entrenamiento.
Carpeta val: almacena los archivos de imagen del conjunto de verificación.
Carpeta de prueba: almacena los archivos de imagen del conjunto de prueba.
Lo anterior es una breve introducción a este script de conversión de formato VOC a COCO. Puede realizar modificaciones y optimizaciones apropiadas según sus necesidades reales. Este script puede ayudarlo a convertir formatos de conjuntos de datos de manera eficiente y admite la copia automática de archivos de imágenes para facilitar la administración y el uso de conjuntos de datos.
Nota: Este artículo y el código son generados de forma completamente automática por GPT 4 y el código probado se puede utilizar normalmente.