Creación de Modelo YOLO con Anotaciones de Unity

Solución en Python para convertir etiquetas JSON a formato YOLO

Highlights

Conversión directa: Transforma las etiquetas JSON de Unity al formato YOLO txt.
Normalización: Calcula coordenadas normalizadas para detección de objetos.
Mapeo de clases: Permite asignar IDs únicos para dados y diferentes piezas de ajedrez.

Introducción y Contexto

En este tutorial, aprenderás a convertir las anotaciones generadas por Unity en formato JSON a un formato compatible con YOLO, el cual es ampliamente utilizado para el entrenamiento de modelos de detección de objetos. Tu conjunto de datos sintéticos consiste en imágenes que contienen dos tipos de objetos: dados y piezas de ajedrez. Además, necesitas que el modelo no solo detecte la presencia de estos objetos, sino que también identifique el tipo específico de pieza de ajedrez o el número que muestra el dado.

Estructura del Archivo JSON y Consideraciones

Generalmente, un archivo JSON de Unity puede tener una estructura similar a la siguiente:


[
  {
    "filename": "imagen1.jpg",
    "width": 640,
    "height": 480,
    "annotations": [
      {
        "class": "dado",
        "number": 5,
        "x_min": 100,
        "y_min": 100,
        "x_max": 150,
        "y_max": 150
      },
      {
        "class": "peon",
        "x_min": 200,
        "y_min": 200,
        "x_max": 230,
        "y_max": 230
      }
    ]
  },
  ...
]

En el ejemplo anterior, cada entrada del JSON contiene detalles relevantes de cada imagen, incluyendo el nombre del archivo, dimensiones y las anotaciones correspondientes a cada objeto. El formato deseado para YOLO requiere que cada imagen tenga un archivo txt donde cada línea siga la sintaxis:

clase x_center y_center ancho alto

Los valores deben estar normalizados entre 0 y 1, es decir, divididos por el ancho y alto de la imagen.

Mapeo de Clases

Para gestionar los diferentes objetos (dados y piezas de ajedrez) es esencial definir un mapeo entre las clases presentes en el JSON y un identificador numérico. La siguiente tabla ilustra un ejemplo de mapeo:

ID	Clase	Descripción
0	dado	Objeto dado, con atributo numérico para identificar su cara
1	peon	Pieza de ajedrez: Peón
2	caballo	Pieza de ajedrez: Caballo
3	torre	Pieza de ajedrez: Torre
4	alfil	Pieza de ajedrez: Alfil
5	reina	Pieza de ajedrez: Reina
6	rey	Pieza de ajedrez: Rey

Puedes modificar o ampliar esta tabla según las piezas de ajedrez que vayas a utilizar y la cantidad de tipos de dados que deseas distinguir.

Script en Python para Conversión a YOLO

A continuación se muestra un código de ejemplo en Python que toma el archivo JSON exportado desde Unity y genera archivos de texto en el formato requerido para YOLO. Este script asume que cada imagen tendrá un archivo .txt con las anotaciones de todos los objetos detectados.

Código Python


# -*- coding: utf-8 -*-
import json
import os

def class_name_to_id(class_name):
    """
    Función para mapear nombres de clase a IDs numéricos.
    Modifica este diccionario según tus clases específicas.
    """
    class_mapping = {
        'dado': 0,
        'peon': 1,
        'caballo': 2,
        'torre': 3,
        'alfil': 4,
        'reina': 5,
        'rey': 6
    }
    return class_mapping.get(class_name, -1)  # Devuelve -1 si no coincide.

def convertir_json_a_yolo(json_path, output_dir):
    """
    Convierte las anotaciones en JSON de Unity al formato YOLO txt.
    Cada imagen tendrá un archivo txt con las anotaciones en formato:
    clase x_center y_center width height
    """
    # Cargar datos JSON
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # Para cada imagen en el JSON
    for item in data:
        filename = item['filename']
        img_width = item['width']
        img_height = item['height']
        annotations = item.get('annotations', [])
        
        # Genera el nombre del archivo txt a partir del nombre de la imagen.
        txt_filename = os.path.splitext(filename)[0] + '.txt'
        txt_filepath = os.path.join(output_dir, txt_filename)
        
        # Abrir el archivo a escribir
        with open(txt_filepath, 'w', encoding='utf-8') as txt_file:
            for obj in annotations:
                class_name = obj.get('class')
                # Obtener el id de la clase
                class_id = class_name_to_id(class_name)
                if class_id == -1:
                    print(f"Clase no reconocida: {class_name}. Se omitirá la anotación.")
                    continue

                # Extraer coordenadas de la bounding box
                # Los nombres de los campos pueden variar; ajusta según tu JSON.
                x_min = obj.get('x_min')
                y_min = obj.get('y_min')
                x_max = obj.get('x_max')
                y_max = obj.get('y_max')

                # Validar que existan las coordenadas
                if None in (x_min, y_min, x_max, y_max):
                    print(f"Coordenadas incompletas para {filename}. Skipping entry.")
                    continue
                
                # Calcular el centro, ancho y alto normalizados
                x_center = (x_min + x_max) / (2.0 * img_width)
                y_center = (y_min + y_max) / (2.0 * img_height)
                width = (x_max - x_min) / img_width
                height = (y_max - y_min) / img_height
                
                # Escribir la anotación en formato YOLO
                txt_file.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
                
                # Si necesitas incluir información extra, por ejemplo el número del dado, puedes agregar comentarios
                if class_name == "dado" and 'number' in obj:
                    txt_file.write(f"# Number: {obj['number']}\n")

    print("Conversión completada.")

if __name__ == "__main__":
    # Ruta al archivo JSON exportado desde Unity
    json_path = "ruta/a/tu/archivo.json"
    # Directorio donde se guardarán los archivos txt de YOLO
    output_dir = "ruta/a/tu/directorio/de/salida"
    
    # Crear el directorio de salida si no existe
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    convertir_json_a_yolo(json_path, output_dir)

Explicación del Código

El script se estructura en distintas secciones:

Mapeo de Clases

La función class_name_to_id realiza un mapeo entre los nombres de los objetos en Unity y sus correspondientes IDs numéricos que utiliza YOLO. Se ha implementado un mecanismo para retornar un valor de -1 en caso de que la clase no coincida, de modo que dicha anotación se ignore.

Conversión de Coordenadas

Por cada objeto anotado en el JSON, se calcula el centro (x_center, y_center) y las dimensiones (width, height) normalizados dividiendo entre el ancho y alto de la imagen.

Guardado de Archivos

Para cada imagen se genera un archivo txt cuyo nombre se deriva del nombre de la imagen. Cada línea en estos archivos sigue el formato requerido por YOLO, asegurando la compatibilidad para el entrenamiento posterior.

Tabla Resumen: Proceso de Conversión

Etapa	Descripción
Lectura del JSON	El script carga el archivo JSON con las anotaciones de cada imagen.
Mapeo de Clases	Se asigna un ID numérico para cada clase (ej. dado, peon, etc.).
Cálculo de Coordenadas	Se normalizan las coordenadas de cada bounding box (x_center, y_center, width, height).
Escritura del Archivo	Se genera un archivo .txt por imagen con las anotaciones en formato YOLO.
Información Extra	Se puede incluir información adicional (como el número en el dado) como comentarios.

Consideraciones Adicionales

Personalización de Estructuras JSON: Es fundamental revisar la estructura exacta del JSON exportado por Unity. Dependiendo de cómo estén definidos los nombres de las claves (por ejemplo, x_min, y_min vs. position y size), puede que necesites ajustar el script para leer correctamente dichas claves.

Validación de Datos: Es recomendable validar que todas las coordenadas se encuentren presentes y que el tamaño de la imagen (ancho y alto) se haya definido correctamente. Si hay inconsistencias en los datos, el script puede omitir ciertas anotaciones o emitir advertencias.

Extensión del Modelo: Dado que tu modelo deberá no solo detectar la presencia de un objeto sino también identificar el número en el dado o el tipo de pieza de ajedrez, tendrás que considerar estrategias adicionales. Una opción es entrenar un modelo multitarea o emplear un segundo modelo de clasificación para analizar las regiones correspondientes a cada objeto.

Próximos Pasos y Entrenamiento del Modelo

Una vez que hayas convertido y validado todas las anotaciones al formato YOLO, el siguiente paso es preparar el conjunto de datos para el entrenamiento del modelo de detección de objetos YOLO. Revisa la documentación de frameworks como Ultralytics para configurar el entorno, ajustar los hiperparámetros y definir el esquema de entrenamiento.

El proceso incluirá la división de los datos en conjuntos de entrenamiento y validación, la configuración de las clases a detectar y, finalmente, la ejecución del entrenamiento del modelo utilizando los archivos txt generados.

Conclusión

En este artículo se ha presentado un enfoque detallado para convertir anotaciones en formato JSON, generadas por Unity, al formato requerido por YOLO. A través de la implementación en Python se abordan los puntos esenciales: el mapeo de clases, la normalización de las coordenadas y la generación de archivos txt para cada imagen. Además, se enfatiza la importancia de adaptar el código a la estructura específica de las anotaciones y de validar que toda la información necesaria esté presente.

Con estos pasos, no solo estarás listo para entrenar un modelo de detección de objetos personalizado, sino que también podrás extender el proceso para tareas más complejas, como la clasificación adicional de números en dados y tipos de piezas de ajedrez. Esta preparación es fundamental para asegurar que el modelo YOLO entienda correctamente la información y realice predicciones precisas en un entorno sintético.