Detección de Anomalías en Imágenes Industriales con Autoencoders Variacionales (VAE) en PyTorch

La detección de anomalías en imágenes es una tarea crítica en muchos campos, desde la medicina hasta la seguridad y, especialmente, en la inspección de calidad industrial. En entornos de fabricación, identificar defectos en productos es fundamental para garantizar estándares y reducir costes.

Tradicionalmente, la inspección se realizaba de forma manual, un proceso propenso a errores humanos y muy lento. La visión artificial ha revolucionado este campo, permitiendo automatizar la detección de defectos con mayor precisión y eficiencia.

En este tutorial, nos centraremos en los Autoencoders Variacionales (VAE), un tipo de red neuronal generativa que se ha mostrado muy efectiva para aprender distribuciones de datos 'normales' y, por ende, identificar desviaciones como anomalías. Utilizaremos PyTorch, una de las librerías más populares para deep learning, para construir y entrenar nuestro modelo.

¿Qué es la Detección de Anomalías en Imágenes? 🖼️

La detección de anomalías se refiere a la identificación de patrones en los datos que no se ajustan al comportamiento esperado. En el contexto de las imágenes industriales, esto significa encontrar imperfecciones, daños o diferencias significativas en un producto que debería ser uniforme.

Existen dos enfoques principales:

1. Supervisado: Se tienen imágenes tanto de productos 'normales' como de productos 'defectuosos' con etiquetas. Esto permite entrenar clasificadores tradicionales.
2. No Supervisado/Semi-supervisado: Solo se dispone de imágenes de productos 'normales'. El modelo aprende a caracterizar lo 'normal' y cualquier desviación significativa se considera una anomalía. Este es el enfoque que usaremos con los VAEs, ya que a menudo es difícil obtener suficientes datos etiquetados de defectos en el mundo real.

Aplicaciones en la Industria 🏭

• Control de calidad en líneas de producción: Detección de rasguños, grietas, decoloraciones, piezas faltantes o mal ensambladas en productos manufacturados.
• Inspección de superficies: Identificación de defectos en materiales como metal, plástico, tela o vidrio.
• Monitoreo de equipos: Detección de desgaste o daño en componentes de maquinaria industrial.

Autoencoders Variacionales (VAE): La Teoría Detrás ✨

Los VAEs son modelos generativos que combinan los principios de los Autoencoders tradicionales con ideas de la inferencia variacional. Un Autoencoder busca aprender una representación latente (codificación) de los datos de entrada en un espacio de menor dimensión y luego reconstruir la entrada a partir de esa representación.

Componentes Principales de un VAE

Un VAE consta de dos partes principales:

1. Encoder (Codificador): Mapea la imagen de entrada x a un espacio latente z. A diferencia de un Autoencoder estándar que produce un punto z, el VAE produce los parámetros de una distribución de probabilidad (generalmente Gaussiana), es decir, la media μ y la varianza σ (o log-varianza log_σ²) para cada dimensión de z.
2. Decoder (Decodificador): Mapea una muestra z del espacio latente de vuelta al espacio de la imagen, produciendo una reconstrucción x'.

La Función de Pérdida del VAE

La función de pérdida de un VAE tiene dos componentes:

1. Pérdida de Reconstrucción (Reconstruction Loss): Mide qué tan bien el decodificador reconstruye la imagen original a partir de la muestra del espacio latente. Comúnmente se usa el Error Cuadrático Medio (MSE) o la Entropía Cruzada Binaria (BCE) para imágenes binarias.
2. Pérdida KL Divergence (KL Loss): Mide la divergencia Kullback-Leibler entre la distribución latente aprendida por el encoder (Q(z|x)) y una distribución a priori (P(z), usualmente una Gaussiana estándar N(0, I)). Esta pérdida incentiva que el espacio latente sea regular y bien estructurado, facilitando la generación de nuevas muestras y la interpolación.

La pérdida total es la suma ponderada de estas dos: L = L_reconstrucción + β * L_KL. El parámetro β se usa para equilibrar la importancia de ambas pérdidas.

Preparación del Entorno 🛠️

Antes de empezar a codificar, necesitamos configurar nuestro entorno de Python. Asegúrate de tener pip instalado.

1. Instalación de Librerías

pip install torch torchvision matplotlib numpy opencv-python scikit-learn

2. Conjunto de Datos

Para este tutorial, utilizaremos un conjunto de datos sintético o uno real simplificado que simule imágenes industriales. Un buen ejemplo podría ser el dataset MVTEC AD, que contiene imágenes de objetos industriales con y sin defectos. Para simplificar, asumiremos que tenemos una carpeta data/train/good con imágenes de productos normales y data/test/good y data/test/anomaly para evaluar.

Estructura de directorios:

. 
├── data/
│   ├── train/
│   │   └── good/
│   │       ├── normal_001.png
│   │       ├── normal_002.png
│   │       └── ...
│   └── test/
│       ├── good/
│       │   ├── normal_test_001.png
│       │   └── ...
│       └── anomaly/
│           ├── defect_001.png
│           └── ...
└── vae_anomaly_detector.py

Implementación del VAE en PyTorch 👨‍💻

Vamos a crear un script Python para nuestro VAE.

1. Importaciones y Configuraciones

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import os
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, auc

# Configuración de dispositivos
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Usando dispositivo: {device}")

# Hiperparámetros
LATENT_DIM = 128
BATCH_SIZE = 32
EPOCHS = 50
LR = 0.001
IMAGE_SIZE = 128 # Todas las imágenes se redimensionarán a 128x128

2. Dataset Personalizado

Necesitamos una clase Dataset para cargar nuestras imágenes.

class AnomalyDetectionDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_paths = []
        for subdir, _, files in os.walk(root_dir):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
                    self.image_paths.append(os.path.join(subdir, file))

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image

# Transformaciones para las imágenes
transform = transforms.Compose([
    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalización estándar de ImageNet
])

# Carga de datos
train_dataset = AnomalyDetectionDataset(root_dir='./data/train/good', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

# Para la evaluación, cargamos tanto las imágenes buenas como las anómalas del conjunto de test
test_good_dataset = AnomalyDetectionDataset(root_dir='./data/test/good', transform=transform)
test_good_loader = DataLoader(test_good_dataset, batch_size=BATCH_SIZE, shuffle=False)

test_anomaly_dataset = AnomalyDetectionDataset(root_dir='./data/test/anomaly', transform=transform)
test_anomaly_loader = DataLoader(test_anomaly_dataset, batch_size=BATCH_SIZE, shuffle=False)

print(f"Imágenes de entrenamiento: {len(train_dataset)}")
print(f"Imágenes de test (buenas): {len(test_good_dataset)}")
print(f"Imágenes de test (anómalas): {len(test_anomaly_dataset)}")

3. Definición del Modelo VAE

Implementaremos un VAE con capas convolucionales. El encoder reducirá progresivamente el tamaño espacial y el decoder lo aumentará.

class VAE(nn.Module):
    def __init__(self, latent_dim=LATENT_DIM, image_size=IMAGE_SIZE):
        super(VAE, self).__init__()

        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=4, stride=2, padding=1), # 128 -> 64
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1), # 64 -> 32
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1), # 32 -> 16
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1), # 16 -> 8
            nn.ReLU(),
            nn.Flatten()
        )
        # Calcular el tamaño de la salida del encoder para la capa lineal
        # Para 128x128, después de 4 Conv2d con stride 2, el tamaño es 128 / (2^4) = 8
        # Canales son 256, por lo que 256 * 8 * 8
        self.fc_mu = nn.Linear(256 * (image_size // 16) * (image_size // 16), latent_dim)
        self.fc_logvar = nn.Linear(256 * (image_size // 16) * (image_size // 16), latent_dim)

        # Decoder
        self.decoder_input = nn.Linear(latent_dim, 256 * (image_size // 16) * (image_size // 16))
        self.decoder = nn.Sequential(
            nn.Unflatten(1, (256, image_size // 16, image_size // 16)),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # 8 -> 16
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 16 -> 32
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1), # 32 -> 64
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1), # 64 -> 128
            nn.Sigmoid() # La salida de la imagen está entre 0 y 1
        )

    def encode(self, x):
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std) # Muestra de una distribución normal estándar
        return mu + eps * std

    def decode(self, z):
        h = self.decoder_input(z)
        x_reconstructed = self.decoder(h)
        return x_reconstructed

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

model = VAE().to(device)
optimizer = optim.Adam(model.parameters(), lr=LR)

4. Función de Pérdida del VAE

La función de pérdida combina el error de reconstrucción (MSE) y la KL Divergence.

def loss_function(recon_x, x, mu, logvar):
    BCE = nn.functional.mse_loss(recon_x, x, reduction='sum') # O BCE si las imágenes son binarias
    # KL Divergence = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

5. Ciclo de Entrenamiento

Entrenaremos el VAE solo con imágenes 'normales'.

def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, data in enumerate(train_loader):
        data = data.to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    print(f'====> Epoch: {epoch} Average loss: {train_loss / len(train_loader.dataset):.4f}')

# Ejecutar el entrenamiento
print("Iniciando entrenamiento...")
for epoch in range(1, EPOCHS + 1):
    train(epoch)
print("Entrenamiento finalizado.")

# Guardar el modelo (opcional)
torch.save(model.state_dict(), "vae_anomaly_detector.pth")
<div class="callout note">📌 <strong>Nota:</strong> Guardar el estado del modelo te permite reutilizarlo sin tener que reentrenar.</div>

Detección de Anomalías y Evaluación 📊

Una vez entrenado el VAE, podemos usarlo para detectar anomalías. La métrica clave será el error de reconstrucción.

1. Función para Calcular el Error de Reconstrucción

def calculate_reconstruction_errors(dataloader):
    model.eval()
    reconstruction_errors = []
    original_images = []
    reconstructed_images = []
    with torch.no_grad():
        for data in dataloader:
            data = data.to(device)
            recon_batch, mu, logvar = model(data)
            # Calculamos el MSE como error de reconstrucción
            error = nn.functional.mse_loss(recon_batch, data, reduction='none').sum(dim=[1,2,3])
            reconstruction_errors.extend(error.cpu().numpy())
            original_images.extend(data.cpu().numpy())
            reconstructed_images.extend(recon_batch.cpu().numpy())
    return np.array(reconstruction_errors), np.array(original_images), np.array(reconstructed_images)

print("Calculando errores de reconstrucción para imágenes de test...")
# Errores para imágenes buenas
good_errors, good_originals, good_recons = calculate_reconstruction_errors(test_good_loader)
# Errores para imágenes anómalas
anomaly_errors, anomaly_originals, anomaly_recons = calculate_reconstruction_errors(test_anomaly_loader)

2. Establecer un Umbral de Anomalía 📈

Para clasificar una imagen como 'anómala' o 'normal', necesitamos un umbral en el error de reconstrucción. Este umbral se suele determinar basándose en los errores de las imágenes normales del conjunto de entrenamiento/validación.

Una estrategia común es calcular el percentil 95 (o 99) de los errores de reconstrucción de las imágenes de entrenamiento normales. Cualquier error por encima de este umbral se considera anómalo.

# Para un umbral más robusto, podríamos calcular esto en un conjunto de validación.
# Para este ejemplo, usaremos el percentil 95 de los errores de las imágenes buenas de test como aproximación.
threshold = np.percentile(good_errors, 95)
print(f"Umbral de anomalía: {threshold:.4f}")

# Visualizar la distribución de errores
plt.figure(figsize=(10, 6))
plt.hist(good_errors, bins=50, alpha=0.5, label='Buenas (Test)', color='blue')
plt.hist(anomaly_errors, bins=50, alpha=0.5, label='Anómalas (Test)', color='red')
plt.axvline(threshold, color='green', linestyle='dashed', linewidth=2, label=f'Umbral: {threshold:.2f}')
plt.xlabel('Error de Reconstrucción (MSE)')
plt.ylabel('Frecuencia')
plt.title('Distribución de Errores de Reconstrucción')
plt.legend()
plt.show()

3. Métricas de Evaluación (ROC AUC) ✅

Para evaluar la efectividad de nuestro detector, podemos usar la curva ROC (Receiver Operating Characteristic) y el área bajo la curva (AUC).

# Combinar errores y crear etiquetas (0 para bueno, 1 para anómalo)
y_true = np.concatenate([np.zeros_like(good_errors), np.ones_like(anomaly_errors)])
y_scores = np.concatenate([good_errors, anomaly_errors])

# Calcular la curva ROC
fpr, tpr, thresholds = roc_curve(y_true, y_scores)
roc_auc = auc(fpr, tpr)

plt.figure(figsize=(8, 8))
plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'Curva ROC (AUC = {roc_auc:.2f})')
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('Tasa de Falsos Positivos (FPR)')
plt.ylabel('Tasa de Verdaderos Positivos (TPR)')
plt.title('Curva ROC para Detección de Anomalías')
plt.legend(loc="lower right")
plt.show()

print(f"Área bajo la curva ROC (AUC): {roc_auc:.4f}")

# Clasificación simple basada en el umbral
y_pred = (y_scores >= threshold).astype(int)

from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

accuracy = accuracy_score(y_true, y_pred)
precision = precision_score(y_true, y_pred)
recall = recall_score(y_true, y_pred)
f1 = f1_score(y_true, y_pred)

print(f"Accuracy: {accuracy:.4f}")
print(f"Precision: {precision:.4f}")
print(f"Recall: {recall:.4f}")
print(f"F1-Score: {f1:.4f}")

4. Visualización de Reconstrucciones y Anomalías 👁️

Es útil visualizar algunas reconstrucciones para entender cómo el VAE maneja las imágenes normales y anómalas.

# Función para desnormalizar imágenes para visualización
def denormalize(tensor):
    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
    return tensor * std + mean

# Visualizar algunas reconstrucciones de imágenes buenas
plt.figure(figsize=(10, 5))
plt.suptitle('Imágenes Buenas: Original vs Reconstrucción')
for i in range(min(5, len(good_originals))):
    # Original
    plt.subplot(2, 5, i + 1)
    img_original = denormalize(torch.tensor(good_originals[i])).permute(1, 2, 0).numpy()
    plt.imshow(img_original)
    plt.title(f'Original\nError: {good_errors[i]:.2f}')
    plt.axis('off')

    # Reconstrucción
    plt.subplot(2, 5, i + 6)
    img_recon = denormalize(torch.tensor(good_recons[i])).permute(1, 2, 0).numpy()
    plt.imshow(img_recon)
    plt.title('Reconstrucción')
    plt.axis('off')
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.show()

# Visualizar algunas reconstrucciones de imágenes anómalas
plt.figure(figsize=(10, 5))
plt.suptitle('Imágenes Anómalas: Original vs Reconstrucción')
for i in range(min(5, len(anomaly_originals))):
    # Original
    plt.subplot(2, 5, i + 1)
    img_original = denormalize(torch.tensor(anomaly_originals[i])).permute(1, 2, 0).numpy()
    plt.imshow(img_original)
    plt.title(f'Original\nError: {anomaly_errors[i]:.2f}')
    plt.axis('off')

    # Reconstrucción
    plt.subplot(2, 5, i + 6)
    img_recon = denormalize(torch.tensor(anomaly_recons[i])).permute(1, 2, 0).numpy()
    plt.imshow(img_recon)
    plt.title('Reconstrucción')
    plt.axis('off')
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.show()

# Opcional: Visualizar mapas de error de reconstrucción
plt.figure(figsize=(10, 5))
plt.suptitle('Mapas de Error de Reconstrucción (Anomalías)')
for i in range(min(5, len(anomaly_originals))):
    original = denormalize(torch.tensor(anomaly_originals[i])).cpu().numpy()
    reconstructed = denormalize(torch.tensor(anomaly_recons[i])).cpu().numpy()
    
    # Calcular el error pixel a pixel
    error_map = np.mean(np.abs(original - reconstructed), axis=0) # Error absoluto promedio por canal
    
    plt.subplot(2, 5, i + 1)
    plt.imshow(original.transpose(1, 2, 0))
    plt.title('Original')
    plt.axis('off')
    
    plt.subplot(2, 5, i + 6)
    plt.imshow(error_map, cmap='hot') # Mapa de calor para el error
    plt.title(f'Error Map\nTotal Error: {anomaly_errors[i]:.2f}')
    plt.axis('off')
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.show()

Optimización y Consideraciones Adicionales 🚀

Hiperparámetros y Arquitectura

• LATENT_DIM: Un espacio latente más grande puede permitir al VAE aprender representaciones más complejas, pero también puede llevar a un sobreajuste si es demasiado grande para la complejidad de los datos.
• EPOCHS y LR: Ajustar el número de épocas y la tasa de aprendizaje es crucial. Un Early Stopping puede ser útil para evitar el sobreentrenamiento.
• Arquitectura del Encoder/Decoder: Experimenta con diferentes números de capas convolucionales, tamaños de kernel, strides y funciones de activación. Modelos más profundos pueden capturar más detalles pero requieren más datos y tiempo de entrenamiento.

Balance β en la Pérdida KL

El β en la función de pérdida BCE + β * KLD es un hiperparámetro importante. Un β más alto regulariza más fuertemente el espacio latente, forzándolo a estar más cerca de la distribución a priori. Esto puede resultar en reconstrucciones de menor fidelidad pero con un espacio latente más 'suave' y generativo. Para detección de anomalías, a menudo se busca un balance para tener buenas reconstrucciones de 'normales' pero también un espacio latente bien comportado.

Detección de Anomalías Localizadas

El error de reconstrucción global es útil para determinar si toda la imagen es anómala. Sin embargo, para anomalías localizadas, es más efectivo analizar el mapa de error pixel a pixel. Al calcular la diferencia absoluta entre la imagen original y la reconstruida, podemos identificar las regiones con mayor discrepancia, lo que a menudo corresponde a la ubicación del defecto.

Consideraciones de Despliegue

• Velocidad: Para aplicaciones en tiempo real, la inferencia debe ser rápida. Optimiza el modelo (por ejemplo, con torch.jit.script o exportación a ONNX) y considera usar hardware acelerador (GPU, TPU).
• Robustez: El modelo debe ser robusto a variaciones de iluminación, oclusiones menores, y otros ruidos no relacionados con anomalías. Aumenta la diversidad del conjunto de datos de entrenamiento con técnicas de data augmentation.

Conclusión 🎉

En este tutorial, hemos cubierto en profundidad cómo implementar un sistema de detección de anomalías en imágenes industriales utilizando Autoencoders Variacionales (VAE) en PyTorch. Hemos explorado los fundamentos teóricos, la preparación de un conjunto de datos, la construcción del modelo, el ciclo de entrenamiento y, lo más importante, cómo evaluar y visualizar los resultados de la detección.

Los VAEs ofrecen una solución potente y flexible para el control de calidad automatizado, permitiendo a las empresas mejorar la eficiencia y la precisión en la identificación de defectos, incluso con una escasez de ejemplos de anomalías. Al dominar estas técnicas, estarás bien equipado para abordar desafíos complejos de visión artificial en el ámbito industrial.

¡Sigue experimentando con diferentes arquitecturas, hiperparámetros y datasets para perfeccionar tus habilidades en la detección de anomalías!