Upload app.py

app.py

@@ -0,0 +1,106 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+import gradio as gr
+
+
+class CNN(nn.Module):
+    def __init__(self, in_channels=1, num_classes=4):
+        super(CNN, self).__init__()
+
+        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        self.batch_norm1 = nn.BatchNorm2d(32)
+        self.batch_norm2 = nn.BatchNorm2d(64)
+        self.batch_norm3 = nn.BatchNorm2d(128)
+
+        self.dropout = nn.Dropout(0.5)
+
+        # Calcular el tamaño de la entrada a la capa fully connected
+        self.fc1 = nn.Linear(128 * (200 // 8) * (200 // 8), 256)
+        self.fc2 = nn.Linear(256, num_classes)
+
+    def forward(self, x):
+        x = F.relu(self.batch_norm1(self.conv1(x)))
+        x = self.pool(x)
+
+        x = F.relu(self.batch_norm2(self.conv2(x)))
+        x = self.pool(x)
+
+        x = F.relu(self.batch_norm3(self.conv3(x)))
+        x = self.pool(x)
+
+        x = x.view(x.shape[0], -1)  # Aplanar
+        x = self.dropout(F.relu(self.fc1(x)))
+        x = self.fc2(x)
+
+        return x
+
+
+model = CNN()
+model.load_state_dict(
+    torch.load("gabriel_complex_modelo.pth", map_location=torch.device("cpu"))
+)
+
+
+def inference(model, imagen, device="cpu"):
+    label_mapping = {0: "Círculo", 1: "Triángulo", 2: "Cuadrado", 3: "Estrella"}
+
+    model.eval()  # Ponemos el modelo en modo evaluación
+
+    # Realizar la inferencia
+    with torch.no_grad():
+        scores = model(imagen)  # Output: tensor con logits
+        probabilities = torch.softmax(
+            scores, dim=1
+        )  # Convertir logits a probabilidades
+        _, prediction = scores.max(1)  # Obtener la clase con mayor probabilidad
+        label_predicho = prediction.item()
+
+    # Diccionario con las probabilidades
+    probabilities_dict = {
+        label_mapping[i]: float(probabilities[0, i]) for i in range(4)
+    }
+
+    return label_mapping[label_predicho], probabilities_dict
+
+
+def predict(img):
+    image_array = img["composite"][:, :, 3]
+    image_array = 255 - image_array
+
+    image_tensor = torch.from_numpy(image_array).unsqueeze(0)
+
+    transform_to_gray = transforms.Compose(
+        [
+            transforms.Resize((200, 200)),
+            transforms.ConvertImageDtype(dtype=torch.float32),  # Convertir a flotante
+        ]
+    )
+
+    image = transform_to_gray(image_tensor)
+    image = image.unsqueeze(0)  # Agregar dimensión extra
+
+    # Hacemos la inferencia
+    label_predict, probabilities = inference(model, image, device="cpu")
+    print(label_predict)
+    print(probabilities)
+
+    return probabilities  # Retorna el diccionario con las probabilidades
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        im = gr.Sketchpad(type="numpy", crop_size="1:1")
+        out = gr.Label()
+
+    im.change(predict, outputs=out, inputs=im, show_progress="hidden")
+
+demo.launch(share=True, debug=False)