Preparación del Conjunto de Datos y Transformaciones

Para el procesamiento de imágenes de resonancia magnética (MRI), es fundamental estructurar los datos de manera eficiente. En lugar de cargar todas las imágenes en la memoria durante la inicialización, utilizaremos la clase Dataset de PyTorch para cargar las imágenes bajo demanda. Esto optimiza el uso de la memoria RAM.

A continuación, se define una clase personalizada que hereda de torch.utils.data.Dataset. Esta implementación utiliza pathlib para la gestión de rutas y PIL para la carga de imágenes, aplicando transformaciones estándar de torchvision.

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from pathlib import Path

class BrainMRIDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = Path(root_dir)
        self.transform = transform
        self.image_paths = []
        self.labels = []

        # Estructura esperada: root_dir/yes/*.jpg y root_dir/no/*.jpg
        for label_name, label_val in [('yes', 1.0), ('no', 0.0)]:
            class_dir = self.root_dir / label_name
            if class_dir.exists():
                for img_path in class_dir.glob('*.jpg'):
                    self.image_paths.append(str(img_path))
                    self.labels.append(label_val)

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

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = self.labels[idx]

        # Cargar imagen y asegurar que esté en formato RGB
        image = Image.open(img_path).convert('RGB')

        if self.transform:
            image = self.transform(image)

        return {
            'image': image, 
            'label': torch.tensor(label, dtype=torch.float32)
        }

Pipeline de Transformaciones y DataLoader

Las imágenes deben ser redimensionadas y convertidas a tensores. El DataLoader se encarga de agrupar los datos en lotes (batches) y aleatorizar el orden para el entrenamiento.

# Definición del pipeline de transformaciones
transform_pipeline = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor()
])

# Instanciar el dataset y el dataloader
dataset = BrainMRIDataset(root_dir='ruta/al/brain_tumor_dataset', transform=transform_pipeline)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=2)

# Verificación de las dimensiones del lote
for batch in dataloader:
    print(f"Dimensiones de las imágenes: {batch['image'].shape}") # torch.Size([32, 3, 128, 128])
    print(f"Dimensiones de las etiquetas: {batch['label'].shape}") # torch.Size([32])
    break

Arquitectura de la Red Neuronal Convloucional (CNN)

Se diseña una arquitectura CNN para la clasificación binaria. El modelo extrae características espaciales mediante capas convolucionales y pooling, seguidas de capas totalmente conectadas para la clasificación final.

import torch.nn as nn
import torch.nn.functional as F

class TumorDetectionCNN(nn.Module):
    def __init__(self):
        super(TumorDetectionCNN, self).__init__()
        
        # Extractor de características
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2), # 128x128 -> 64x64
            
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)  # 64x64 -> 32x32
        )
        
        # Clasificador (32 * 32 * 32 = 32768)
        self.classifier = nn.Sequential(
            nn.Linear(in_features=32768, out_features=128),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=128, out_features=1)
        )

    def forward(self, x):
        features = self.feature_extractor(x)
        flattened = features.view(features.size(0), -1)
        logits = self.classifier(flattened)
        return torch.sigmoid(logits)

Ciclo de Entrenamiento y Evaluación

El entrenamiento utiliza la pérdida de entropía cruzada binaria (BCELoss) y el optimizador Adam. Se implementa un bucle de entrenamiento estándar y una fase de evaluación para calcular la precisión del modelo.

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = TumorDetectionCNN().to(device)

criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# --- Fase de Entrenamiento ---
num_epochs = 15
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    
    for batch in dataloader:
        inputs = batch['image'].to(device)
        # Ajustar dimensiones de las etiquetas para que coincidan con la salida [batch_size, 1]
        labels = batch['label'].to(device).unsqueeze(1) 

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)

    epoch_loss = running_loss / len(dataset)
    print(f'Epoch {epoch+1}/{num_epochs}, Pérdida: {epoch_loss:.4f}')

# --- Fase de Evaluación ---
model.eval()
correct_preds = 0
total_samples = 0

with torch.no_grad():
    for batch in dataloader:
        inputs = batch['image'].to(device)
        labels = batch['label'].to(device)

        outputs = model(inputs)
        # Aplicar umbral de 0.5 para clasificación binaria
        predicted = (outputs >= 0.5).float().squeeze()

        correct_preds += (predicted == labels).sum().item()
        total_samples += labels.size(0)

accuracy = 100 * correct_preds / total_samples
print(f'Precisión en el conjunto de evaluación: {accuracy:.2f}%')