Creación de un Pipeline con YOLO y LayoutLMv3

Author

Aldo Jaenes

Butler Scientifics Logo

Introducción

En este documento, aprenderás a crear un pipeline que integra YOLO y LayoutLMv3, y cómo entrenarlo utilizando LabelImg para el etiquetado de datos. Este pipeline te permitirá detectar objetos en imágenes (con YOLO) y luego procesar información estructurada de documentos (con LayoutLMv3).

1. Visión General del Pipeline

  • Etiquetado de Datos: Utilizar LabelImg para anotar las regiones de interés en tus imágenes/documentos.
  • Entrenamiento de YOLO: Entrenar un modelo YOLO para la detección de objetos específicos en tus imágenes.
  • Procesamiento con LayoutLMv3: Utilizar las salidas de YOLO (regiones detectadas) como entradas para LayoutLMv3, permitiendo el reconocimiento y comprensión del contenido dentro de esas regiones.
  • Integración: Combinar ambos modelos en un pipeline que tome una imagen de entrada, detecte objetos y procese la información textual y estructural.

2. Configuración del Entorno

Requisitos Previos

  • Python 3.7+
  • PyTorch
  • Transformers de Hugging Face
  • LabelImg
  • OpenCV

Instalación de Paquetes Necesarios

Ejecuta los siguientes comandos en tu terminal para instalar las dependencias necesarias:

pip install torch torchvision torchaudio
pip install transformers
pip install opencv-python
pip install labelImg

3. Etiquetado de Datos con LabelImg

Instalación de LabelImg

Puedes instalar LabelImg usando pip:

pip install labelImg

O clonando el repositorio:

git clone https://github.com/tzutalin/labelImg.git
cd labelImg
python setup.py install

Etiquetado de Imágenes

  1. Abrir LabelImg:

    labelImg

  2. Crear Etiquetas: Define las clases que quieres detectar (por ejemplo, “tabla”, “gráfico”, “texto”).

  3. Anotar Imágenes: Dibuja cajas delimitadoras alrededor de los objetos de interés en tus imágenes.

  4. Guardar Anotaciones: Guarda las anotaciones en formato Pascal VOC (XML) o YOLO (TXT), dependiendo de las necesidades.

4. Preparación de Datos para YOLO

Estructura de Carpetas

Organiza tus datos de la siguiente manera:

dataset/
├── images/
│   ├── train/
│   ├── val/
└── labels/
    ├── train/
    ├── val/
  • images/train/ y images/val/: Contienen las imágenes de entrenamiento y validación.
  • labels/train/ y labels/val/: Contienen las etiquetas correspondientes.

Conversión de Anotaciones (si es necesario)

Si tus anotaciones están en formato Pascal VOC, podrías necesitar convertirlas al formato YOLO. Existen scripts disponibles para esta conversión.

5. Entrenamiento de YOLO

Selección de la Versión de YOLO

Puedes utilizar YOLOv5 o YOLOv8. Para este ejemplo, usaremos YOLOv5.

Clonar el Repositorio de YOLOv5

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

Configuración del Archivo de Datos

Crea un archivo YAML (por ejemplo, custom_data.yaml) con la siguiente estructura:

train: /ruta/a/dataset/images/train
val: /ruta/a/dataset/images/val

nc: Número_de_clases
names: ['clase1', 'clase2', ...]

Entrenamiento

Ejecuta el script de entrenamiento:

python train.py --img 640 --batch 16 --epochs 50 --data custom_data.yaml --weights yolov5s.pt
  • –img: Tamaño de las imágenes.
  • –batch: Tamaño del batch.
  • –epochs: Número de épocas.
  • –data: Ruta al archivo de datos YAML.
  • –weights: Pesos pre-entrenados a utilizar.

6. Uso de YOLO para la Detección

Después del entrenamiento, puedes utilizar el modelo para detectar objetos en nuevas imágenes:

from yolov5.models.common import DetectMultiBackend
from yolov5.utils.torch_utils import select_device
from yolov5.utils.augmentations import letterbox
import cv2
import torch

# Configuración
device = select_device('')
model = DetectMultiBackend('ruta/al/modelo/entrenado.pt', device=device)

7. Preparación para LayoutLMv3

Extracción de Información de las Detecciones de YOLO

Para cada detección de YOLO:

  1. Obtener las Coordenadas de las Cajas Delimitadoras.
  2. Recortar las Regiones de Interés (opcional, si necesitas procesar solo esas regiones).
  3. Extraer el Texto dentro de esas Regiones (puedes utilizar OCR como Tesseract o herramientas de OCR más avanzadas).

Instalación de Dependencias para LayoutLMv3

pip install datasets seqeval

8. Entrenamiento y Uso de LayoutLMv3

Preparación de Datos

LayoutLMv3 requiere entradas específicas:

  • Imagen: La imagen del documento.
  • Texto: Lista de tokens.
  • Caja: Coordenadas de los tokens en la página.
  • Etiquetas: Etiquetas correspondientes a cada token.

Puedes utilizar la librería Datasets de Hugging Face para crear un dataset compatible.

Ejemplo de Uso

from transformers import LayoutLMv3ForTokenClassification, LayoutLMv3Processor
import torch

# Cargar el modelo y el procesador
model = LayoutLMv3ForTokenClassification.from_pretrained('microsoft/layoutlmv3-base', num_labels=Num_Clases)
processor = LayoutLMv3Processor.from_pretrained('microsoft/layoutlmv3-base', apply_ocr=False)

# Procesamiento de la imagen y el texto
encoding = processor(image, text, boxes=boxes, word_labels=labels, return_tensors="pt")

# Inferencia
outputs = model(**encoding)
logits = outputs.logits

Entrenamiento Personalizado

Para entrenar el modelo con tus propios datos:

  1. Preparar el Dataset: Asegúrate de que tus datos estén en el formato correcto.

  2. Definir el Entrenador: Utiliza la clase Trainer de Hugging Face.

    from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    evaluation_strategy="steps",
    save_steps=10,
    eval_steps=10,
    logging_dir='./logs',
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=processor.feature_extractor,
)

trainer.train()

9. Integración del Pipeline Completo

Puedes crear un script que:

  1. Tome una imagen de entrada.
  2. Utilice el modelo YOLO entrenado para detectar objetos.
  3. Extraiga las regiones de interés y el texto dentro de ellas.
  4. Procese estas regiones con LayoutLMv3 para tareas de comprensión.

Ejemplo Simplificado

# Paso 1: Detección con YOLO
detections = model_yolo(image)

# Paso 2: Procesamiento de cada detección
for detection in detections:
    bbox = detection['bbox']
    cropped_image = image.crop(bbox)
    
    # Extracción de texto (usando OCR)
    text = ocr_function(cropped_image)
    
    # Preparación para LayoutLMv3
    encoding = processor(cropped_image, text, boxes=bbox_coords, return_tensors="pt")
    
    # Inferencia con LayoutLMv3
    outputs = model_layoutlmv3(**encoding)
    # Procesamiento de los resultados...

10. Consideraciones Finales

  • OCR: La precisión del OCR es crucial. Considera utilizar herramientas avanzadas como Tesseract con modelos entrenados o servicios como Google Cloud Vision.
  • Ajuste de Hiperparámetros: Experimenta con diferentes parámetros durante el entrenamiento para mejorar el rendimiento.
  • Evaluación: Utiliza métricas como F1-score, precisión y recall para evaluar tus modelos.

Recursos Adicionales