Carga de los datos

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transformaciones = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])  # imágenes en escala de grises
])

dataset = datasets.ImageFolder('ruta/del/dataset', transform=transformaciones)
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size

from torch.utils.data import random_split
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

Arquitectura de la red neuronal

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

class RedCNN(nn.Module):
    def __init__(self):
        super(RedCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64 * 16 * 16, 128)
        self.fc2 = nn.Linear(128, 2)  # 2 clases: fracturado / no fracturado

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # 128x128 → 64x64
        x = self.pool(F.relu(self.conv2(x)))  # 64x64 → 32x32
        x = self.pool(F.relu(self.conv3(x)))  # 32x32 → 16x16
        x = x.view(-1, 64 * 16 * 16)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

Entrenamiento del modelo

import torch.optim as optim

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
modelo = RedCNN().to(device)

criterio = nn.CrossEntropyLoss()
optimizador = optim.Adam(modelo.parameters(), lr=0.001)

for epoca in range(5):
    modelo.train()
    perdida_total = 0
    for imagenes, etiquetas in train_loader:
        imagenes, etiquetas = imagenes.to(device), etiquetas.to(device)

        optimizador.zero_grad()
        salida = modelo(imagenes)
        perdida = criterio(salida, etiquetas)
        perdida.backward()
        optimizador.step()

        perdida_total += perdida.item()

    print(f"Época {epoca+1}, pérdida: {perdida_total:.4f}")

Evaluación del modelo

from sklearn.metrics import confusion_matrix, classification_report
import numpy as np

modelo.eval()
todos_los_resultados = []
todas_las_etiquetas = []

with torch.no_grad():
    for imagenes, etiquetas in test_loader:
        imagenes = imagenes.to(device)
        salida = modelo(imagenes)
        _, predicciones = torch.max(salida, 1)
        todos_los_resultados.extend(predicciones.cpu().numpy())
        todas_las_etiquetas.extend(etiquetas.numpy())

print("Reporte de clasificación:\n")
print(classification_report(todas_las_etiquetas, todos_los_resultados))

matriz = confusion_matrix(todas_las_etiquetas, todos_los_resultados)
print("Matriz de confusión:\n", matriz)

Visualización de resultados

import matplotlib.pyplot as plt
import torchvision

def mostrar_imagenes(imagenes, etiquetas, predicciones=None):
    imagenes = imagenes / 2 + 0.5  # desnormalizar
    npimg = torchvision.utils.make_grid(imagenes).numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    if predicciones is not None:
        plt.title(f"Predicciones: {predicciones}")
    else:
        plt.title(f"Etiquetas reales: {etiquetas}")
    plt.show()

Gráfico de pérdida por época

import matplotlib.pyplot as plt

# Guarda las pérdidas en cada época durante el entrenamiento
perdidas_epocas = []

for epoca in range(5):
    modelo.train()
    perdida_total = 0
    for imagenes, etiquetas in train_loader:
        imagenes, etiquetas = imagenes.to(device), etiquetas.to(device)

        optimizador.zero_grad()
        salida = modelo(imagenes)
        perdida = criterio(salida, etiquetas)
        perdida.backward()
        optimizador.step()

        perdida_total += perdida.item()

    perdidas_epocas.append(perdida_total)
    print(f"Época {epoca+1}, pérdida: {perdida_total:.4f}")

# Gráfico de la pérdida
plt.plot(range(1, len(perdidas_epocas)+1), perdidas_epocas, marker='o')
plt.title("Pérdida durante el entrenamiento")
plt.xlabel("Época")
plt.ylabel("Pérdida")
plt.grid(True)
plt.show()

Matriz de confusión visual

import seaborn as sns
from sklearn.metrics import confusion_matrix

cm = confusion_matrix(todas_las_etiquetas, todos_los_resultados)

plt.figure(figsize=(5, 4))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=['No fracturado', 'Fracturado'], yticklabels=['No fracturado', 'Fracturado'])
plt.ylabel('Etiqueta real')
plt.xlabel('Predicción')
plt.title('Matriz de Confusión')
plt.show()

Mostrar imágenes con sus predicciones

import torchvision
import numpy as np

# Tomamos un lote del conjunto de prueba
dataiter = iter(test_loader)
imagenes, etiquetas = next(dataiter)

imagenes = imagenes.to(device)
salida = modelo(imagenes)
_, predicciones = torch.max(salida, 1)

# Desnormalizamos las imágenes para mostrarlas bien
imagenes = imagenes.cpu() * 0.5 + 0.5

# Mostrar las primeras 6 imágenes con predicción y etiqueta real
fig = plt.figure(figsize=(12, 6))
for idx in range(6):
    ax = fig.add_subplot(2, 3, idx+1, xticks=[], yticks=[])
    img = np.transpose(imagenes[idx].numpy(), (1, 2, 0))
    ax.imshow(img)
    ax.set_title(f"Real: {etiquetas[idx]}\nPred: {predicciones[idx].item()}")
plt.tight_layout()
plt.show()

Guardar los gráficos como imágenes

plt.savefig("nombre_del_grafico.png", dpi=300)