Update app.py

app.py

@@ -5,20 +5,18 @@ import numpy as np
 from PIL import Image
 import torchvision.transforms as transforms
 
-# Modelo autoencoder
+# Modelo Autoencoder
 class Autoencoder(nn.Module):
     def __init__(self):
         super(Autoencoder, self).__init__()
-        # Encoder
-        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1)  # 64x64 -> 32x32
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)  # 32x32 -> 16x16
-        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)  # 16x16 -> 8x8
-        self.fc1 = nn.Linear(128 * 8 * 8, 32)  # Espacio latente
-        # Decoder
+        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
+        self.fc1 = nn.Linear(128 * 8 * 8, 32)
         self.fc2 = nn.Linear(32, 128 * 8 * 8)
-        self.conv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1)  # 8x8 -> 16x16
-        self.conv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1)  # 16x16 -> 32x32
-        self.conv6 = nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1)  # 32x32 -> 64x64
+        self.conv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.conv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.conv6 = nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1)
 
     def encode(self, x):
         z = torch.tanh(self.conv1(x))
@@ -44,29 +42,35 @@ model = Autoencoder()
 model.load_state_dict(torch.load("autoencoder.pth", map_location=torch.device("cpu")))
 model.eval()
 
-# Transformación de entrada
+# Transformación
 transform = transforms.Compose([
     transforms.Grayscale(),
     transforms.Resize((64, 64)),
     transforms.ToTensor()
 ])
 
-# Función de inferencia
+# Umbral de error (ajustable)
+THRESHOLD = 0.01
+
+# Función de predicción
 def detectar_anomalia(imagen):
+    img_tensor = transform(imagen).unsqueeze(0)
     with torch.no_grad():
-        img_tensor = transform(imagen).unsqueeze(0)  # Añadir batch
-        reconstruida = model(img_tensor).squeeze(0).squeeze(0)
-        return reconstruida.numpy()  # Convertir a numpy para visualización
+        reconstruida = model(img_tensor)
 
+    mse = torch.mean((img_tensor - reconstruida) ** 2).item()
+    resultado = "Anómala" if mse > THRESHOLD else "Normal"
+    return resultado
 
-# Interfaz de Gradio
-interface = gr.Interface(
+# Interfaz Gradio
+demo = gr.Interface(
     fn=detectar_anomalia,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="numpy"), gr.Text()],
-    title="Detección de Anomalías con Autoencoder",
-    description="Sube una imagen para detectar anomalías usando un autoencoder entrenado."
+    inputs=gr.Image(type="pil", label="Sube una imagen para analizar"),
+    outputs=gr.Label(label="Resultado"),
+    examples=["anomalous.png", "normal.png"],
+    title="📌 Detección de Anomalías con Autoencoder (PyTorch)",
+    description="Este Space utiliza un autoencoder entrenado con PyTorch para detectar anomalías en imágenes de textiles industriales.",
 )
 
-interface.launch()
+demo.launch()