Objetivos del proyecto
- Segmentar píxel a píxel deslizamientos de tierra mediante una arquitectura UNet.
- Implementar el flujo de entrenamiento con PyTorch.
- Crear un sistema de detección y visualización de deslizamientos.
- Capacidad de inferencia en imágenes individuales, videos o lotes de imágenes de sensores remotos.
1. Configuración del entorno
1.1. Instalación de CUDA
Verifique que su GPU tenga el controlador NVIDIA adecuado y que soporte CUDA 11.8:
nvidia-smi
1.2. Creación del entorno virtual con Anaconda
conda create --name unet-seg python=3.9
conda activate unet-seg
1.3. Instalación de dependencais
pip install torch torchvision torchaudio
pip install matplotlib numpy opencv-python tqdm scikit-learn
pip install albumentations # aumento de datos (opcional)
pip install tensorboard # visualización del entrenamiento (opcional)
2. Preparación del conjunto de datos (formato YOLO a máscaras)
UNet requiere pares imagen-máscara. Las anotaciones en formato YOLO (archivos .txt) deben convertirse a imágenes de máscara en escala de grises donde cada píxel es 0 (fondo) o 1 (deslizamiento).
Script de conversión: convertir_yolo_a_mascara.py
import os
import cv2
import numpy as np
TAMANO_IMAGEN = (512, 512)
def yolo_a_mascara(ruta_txt, tamano):
alto, ancho = tamano
mascara = np.zeros((alto, ancho), dtype=np.uint8)
with open(ruta_txt, 'r') as f:
lineas = f.readlines()
for linea in lineas:
datos = list(map(float, linea.strip().split()))
id_clase = int(datos[0])
coord = np.array(datos[1:]).reshape(-1, 2)
coord[:, 0] *= ancho
coord[:, 1] *= alto
coord = np.round(coord).astype(np.int32)
cv2.fillPoly(mascara, [coord], color=1) # clase = 1
return mascara
# Rutas
dir_img = "dataset/images/train"
dir_etiqueta = "dataset/labels/train"
dir_salida = "dataset/masks/train"
os.makedirs(dir_salida, exist_ok=True)
for archivo_etiqueta in os.listdir(dir_etiqueta):
nombre_base = os.path.splitext(archivo_etiqueta)[0]
ruta_img = os.path.join(dir_img, nombre_base + ".jpg")
ruta_etiqueta = os.path.join(dir_etiqueta, archivo_etiqueta)
if not os.path.exists(ruta_img):
continue
mascara = yolo_a_mascara(ruta_etiqueta, TAMANO_IMAGEN)
cv2.imwrite(os.path.join(dir_salida, nombre_base + ".png"), mascara * 255)
La estructura final del dataset será:
dataset/
├── images/
│ └── train/
└── masks/
└── train/
3. Definición de la clase Dataset y del modelo UNet
3.1. Clase personalizada DatasetDeslizamiento
import os
from torch.utils.data import Dataset
from PIL import Image
import numpy as np
import torch
import torchvision.transforms as transforms
class DatasetDeslizamiento(Dataset):
def __init__(self, dir_img, dir_mask, transform=None):
self.dir_img = dir_img
self.dir_mask = dir_mask
self.transform = transform
self.imagenes = os.listdir(dir_img)
def __len__(self):
return len(self.imagenes)
def __getitem__(self, idx):
ruta_img = os.path.join(self.dir_img, self.imagenes[idx])
ruta_mask = os.path.join(self.dir_mask, self.imagenes[idx].replace('.jpg', '.png'))
img = np.array(Image.open(ruta_img).convert("RGB"))
mask = np.array(Image.open(ruta_mask).convert("L"))
# Binarizar máscara: fondo=0, deslizamiento=1
mask[mask > 0] = 1
if self.transform:
aumentos = self.transform(image=img, mask=mask)
img = aumentos["image"]
mask = aumentos["mask"]
img = transforms.ToTensor()(img)
mask = torch.from_numpy(mask).long()
return img, mask
3.2. Modelo UNet simplificado
import torch.nn as nn
import torch.nn.functional as F
class UNet(nn.Module):
def __init__(self, canales_entrada=3, canales_salida=1):
super(UNet, self).__init__()
def capa_conv_bn_relu(entrada, salida):
return nn.Sequential(
nn.Conv2d(entrada, salida, kernel_size=3, padding=1),
nn.BatchNorm2d(salida),
nn.ReLU(inplace=True),
nn.Conv2d(salida, salida, kernel_size=3, padding=1),
nn.BatchNorm2d(salida),
nn.ReLU(inplace=True)
)
self.enc1 = capa_conv_bn_relu(canales_entrada, 64)
self.enc2 = capa_conv_bn_relu(64, 128)
self.enc3 = capa_conv_bn_relu(128, 256)
self.enc4 = capa_conv_bn_relu(256, 512)
self.bottleneck = capa_conv_bn_relu(512, 1024)
self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
self.dec4 = capa_conv_bn_relu(1024, 512)
self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
self.dec3 = capa_conv_bn_relu(512, 256)
self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
self.dec2 = capa_conv_bn_relu(256, 128)
self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
self.dec1 = capa_conv_bn_relu(128, 64)
self.conv_final = nn.Conv2d(64, canales_salida, kernel_size=1)
def forward(self, x):
enc1 = self.enc1(x)
enc2 = self.enc2(F.max_pool2d(enc1, kernel_size=2, stride=2))
enc3 = self.enc3(F.max_pool2d(enc2, kernel_size=2, stride=2))
enc4 = self.enc4(F.max_pool2d(enc3, kernel_size=2, stride=2))
bottleneck = self.bottleneck(F.max_pool2d(enc4, kernel_size=2, stride=2))
dec4 = self.upconv4(bottleneck)
dec4 = torch.cat((dec4, enc4), dim=1)
dec4 = self.dec4(dec4)
dec3 = self.upconv3(dec4)
dec3 = torch.cat((dec3, enc3), dim=1)
dec3 = self.dec3(dec3)
dec2 = self.upconv2(dec3)
dec2 = torch.cat((dec2, enc2), dim=1)
dec2 = self.dec2(dec2)
dec1 = self.upconv1(dec2)
dec1 = torch.cat((dec1, enc1), dim=1)
dec1 = self.dec1(dec1)
return torch.sigmoid(self.conv_final(dec1))
4. Entrenamiento del modelo
Script completo de entrenamiento
from torch.utils.data import DataLoader
import torch.optim as optim
import torch.nn as nn
from tqdm import tqdm
# Hiperparámetros
EPOCHS = 20
TAM_LOTE = 8
TASA_APRENDIZAJE = 1e-4
DISPOSITIVO = "cuda" if torch.cuda.is_available() else "cpu"
# Cargar datasets
train_set = DatasetDeslizamiento(
dir_img="dataset/images/train",
dir_mask="dataset/masks/train"
)
val_set = DatasetDeslizamiento(
dir_img="dataset/images/val",
dir_mask="dataset/masks/val"
)
train_loader = DataLoader(train_set, batch_size=TAM_LOTE, shuffle=True)
val_loader = DataLoader(val_set, batch_size=TAM_LOTE)
# Inicializar modelo, optimizador y pérdida
modelo = UNet(canales_entrada=3, canales_salida=1).to(DISPOSITIVO)
optimizador = optim.Adam(modelo.parameters(), lr=TASA_APRENDIZAJE)
criterio = nn.BCELoss() # pérdida de entropía cruzada binaria
# Bucle de entrenamiento
for epoca in range(EPOCHS):
modelo.train()
barra = tqdm(train_loader, total=len(train_loader))
for imagenes, mascaras in barra:
imagenes = imagenes.to(DISPOSITIVO)
mascaras = mascaras.float().unsqueeze(1).to(DISPOSITIVO)
predicciones = modelo(imagenes)
perdida = criterio(predicciones, mascaras)
optimizador.zero_grad()
perdida.backward()
optimizador.step()
barra.set_postfix(loss=perdida.item())
# Validación
modelo.eval()
with torch.no_grad():
perdida_val = 0
for imagenes, mascaras in val_loader:
imagenes = imagenes.to(DISPOSITIVO)
mascaras = mascaras.float().unsqueeze(1).to(DISPOSITIVO)
preds = modelo(imagenes)
perdida_val += criterio(preds, mascaras).item()
print(f"Época {epoca+1} | Pérdida validación: {perdida_val/len(val_loader):.4f}")
# Guardar el modelo entrenado
torch.save(modelo.state_dict(), "unet_deslizamiento.pth")
5. Inferencia y sistema de detección
Función de inferencia para una sola imagen
def predecir_imagen(modelo, ruta_imagen, dispositivo="cuda"):
modelo.eval()
img = Image.open(ruta_imagen).convert("RGB")
tensor_img = transforms.ToTensor()(img).unsqueeze(0).to(dispositivo)
with torch.no_grad():
salida = modelo(tensor_img)
mascara_pred = (salida.squeeze().cpu().numpy() > 0.5).astype(np.uint8)
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(img)
plt.title("Imagen original")
plt.subplot(1, 2, 2)
plt.imshow(mascara_pred, cmap='gray')
plt.title("Máscara predicha")
plt.show()