app.py

import torch
import torch.nn as nn
from torchvision import models
import gradio as gr
import os
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from safetensors.torch import load_model
from datasets import load_dataset


# modelos
class Stem(nn.Module):
    def __init__(self):
        super(Stem, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels // 4, stride=1, kernel_size=1),
            nn.BatchNorm2d(out_channels // 4),
            nn.ReLU(inplace=True),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                out_channels // 4,
                out_channels // 4,
                stride=stride,
                kernel_size=3,
                padding=1,
            ),
            nn.BatchNorm2d(out_channels // 4),
            nn.ReLU(inplace=True),
        )

        self.conv3 = nn.Sequential(
            nn.Conv2d(out_channels // 4, out_channels, kernel_size=1, stride=1),
            nn.BatchNorm2d(out_channels),
        )

        self.shortcut = (
            nn.Identity()
            if in_channels == out_channels
            else nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
                nn.BatchNorm2d(out_channels),
            )
        )

        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        identity = self.shortcut(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x += identity
        x = self.relu(x)
        return x


def make_layer(in_channels, out_channels, block, num_blocks):
    layers = []
    for i in range(num_blocks):
        layers.append(block(in_channels, out_channels))
        in_channels = out_channels

    return layers


class FromZero(nn.Module):
    def __init__(self, num_classes=10):
        super(FromZero, self).__init__()
        self.stem = Stem()
        self.layer1 = nn.Sequential(*make_layer(64, 64, ResidualBlock, 2))
        self.layer2 = nn.Sequential(
            ResidualBlock(64, 128, stride=2), ResidualBlock(128, 128)
        )
        self.layer3 = nn.Sequential(
            ResidualBlock(128, 256, stride=2), ResidualBlock(256, 256)
        )
        self.layer4 = nn.Sequential(
            ResidualBlock(256, 512, stride=2), ResidualBlock(512, 512)
        )

        self.flatten = nn.Flatten()
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes)

    def forward(self, x):
        x = self.stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = self.flatten(x)
        x = self.fc(x)
        return x


class PreTrained(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.model = models.resnet18(
            weights=models.ResNet18_Weights.IMAGENET1K_V1, progress=True
        )
        for param in self.model.parameters():
            param.requires_grad = False

        self.model.fc = nn.Sequential(
            nn.Linear(self.model.fc.in_features, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, num_classes),
        )

    def forward(self, x):
        return self.model(x)


with open("etiquetas.txt", "r") as f:
    etiquetas = f.read().splitlines()[1:]
num_clases = len(etiquetas)
codigo = {etiqueta.lower(): i for i, etiqueta in enumerate(etiquetas)}


def codificar_etiqueta(etiqueta):
    return codigo[etiqueta]

dataset = load_dataset(
    "minoruskore/elementosparaevaluarclases", split="train"
)


class imagenDataset(Dataset):
    def __init__(self, dt, transform):
        self.dt = dt
        self.tr = transform

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

    def __getitem__(self, idx):
        row = self.dt[idx]
        imagen = row["image"].convert("RGB")
        label = row["etiqueta"].lower()
        label = codificar_etiqueta(label)
        imagen = self.tr(imagen)
        return imagen, label


tr = transforms.Compose([transforms.Resize([256, 256]), transforms.ToTensor()])
test_dataset = imagenDataset(dataset, transform=tr)
cpus = os.cpu_count()
test_dataloader = DataLoader(test_dataset, batch_size=500, num_workers=cpus)


def multiclass_accuracy(predictions, labels):

    # Obtén las clases predichas (la clase con la mayor probabilidad)
    _, predicted_classes = torch.max(predictions, 1)

    # Compara las clases predichas con las etiquetas verdaderas
    correct_predictions = (predicted_classes == labels).sum().item()

    # Calcula la precisión
    accuracy = correct_predictions / labels.size(0)

    return accuracy


def cargar_evaluar_modelo(archivo, tipo_modelo):
    try:
        if tipo_modelo == "tarea_7":
            modelo = FromZero(num_clases)

        elif tipo_modelo == "tarea_8":
            modelo = PreTrained(num_clases)

        load_model(modelo, archivo)
        modelo.eval()
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        modelo.to(device)
        accuracy = 0
        with torch.no_grad():
            for imagenes, etiquetas in test_dataloader:
                imagenes = imagenes.to(device)
                etiquetas = etiquetas.to(device)
                predictions = modelo(imagenes)
                accuracy += multiclass_accuracy(predictions, etiquetas)
        accuracy = accuracy / len(test_dataloader)
        return accuracy
    except Exception as e:
        return f"Error: {str(e)}"


def evaluate_interface(model_file, model_type):
    if model_file is None:
        return "Por favor, carga un archivo .safetensor"

    # Verificamos que el archivo sea .safetensor
    if not model_file.name.endswith(".safetensor"):
        return "Por favor, carga un archivo con extensión .safetensor"

    # Evaluamos el modelo
    accuracy = cargar_evaluar_modelo(
        model_file.name,
        model_type,
    )

    if isinstance(accuracy, float):
        return f"Precisión del modelo: {accuracy*100:.2f}%"
    else:
        return accuracy


demo = gr.Interface(
    fn=evaluate_interface,
    inputs=[
        gr.File(label="Archivo del modelo (.safetensor)"),
        gr.Radio(["tarea_7", "tarea_8"], label="Tipo de modelo", value="tarea_7"),
    ],
    outputs=gr.Textbox(label="Resultado", lines=1),
    title="Evaluador de Tareas 7 y 8",
    description="Carga un archivo .safetensor de la tarea 7 o 8 y evalúa su precisión en el conjunto de datos de evaluación.",
)

demo.launch()