Update app.py

app.py

@@ -3,13 +3,14 @@ import torch.nn.functional as F
 from torch_geometric.nn import GCNConv, global_mean_pool
 from torch_geometric.utils import from_networkx
 from torch_geometric.data import Data
+
 import networkx as nx
 import gradio as gr
 import matplotlib.pyplot as plt
+from PIL import Image
 import io
-import base64
 
-# ---------- MODELO GCN ----------
+# ---------- MODELO ----------
 
 class GCN(torch.nn.Module):
     def __init__(self, hidden_channels=64):
@@ -33,38 +34,32 @@ model = GCN().to(device)
 model.load_state_dict(torch.load("model_gcn.pth", map_location=device))
 model.eval()
 
-# ---------- FUNCIÓN PARA PARSEAR INPUT ----------
+# ---------- FUNCIONES AUXILIARES ----------
 
 def parse_input(num_nodes, edges_str, node_features_str):
     G = nx.Graph()
 
     try:
-        # Añadimos nodos
+        # Añadir nodos
         for i in range(num_nodes):
             G.add_node(i)
 
-        # Parseamos aristas
-        edges = eval(edges_str)  # formato esperado: [(0,1), (1,2)]
+        # Parsear aristas
+        edges = eval(edges_str)
         G.add_edges_from(edges)
 
-        # Parseamos características
-        node_features = eval(node_features_str)  # formato: [[1,0,1,0,1,0,1], [...]]
+        # Parsear características
+        node_features = eval(node_features_str)
         if len(node_features) != num_nodes or any(len(f) != 7 for f in node_features):
-            raise ValueError("Las características deben ser listas de longitud 7 para cada nodo.")
+            raise ValueError("Cada nodo debe tener exactamente 7 características.")
 
         nx.set_node_attributes(G, {i: node_features[i] for i in range(num_nodes)}, "x")
-
         return G
 
     except Exception as e:
-        raise gr.Error(f"Error al procesar el input: {e}")
-
-# ---------- VISUALIZACIÓN ----------
+        raise gr.Error(f"Error en los datos del grafo: {e}")
 
 def draw_graph(G, pred_label):
-    import matplotlib.pyplot as plt
-    import io
-
     pos = nx.spring_layout(G)
     node_colors = ['lightgreen' if pred_label == 1 else 'lightcoral'] * G.number_of_nodes()
 
@@ -74,12 +69,9 @@ def draw_graph(G, pred_label):
 
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
-    buf.seek(0)
     plt.close()
-
-    return buf  # Devuelve un objeto tipo archivo
-
-# ---------- FUNCIÓN DE PREDICCIÓN ----------
+    buf.seek(0)
+    return Image.open(buf)
 
 def predict_graph(num_nodes, edges_str, node_features_str):
     G = parse_input(num_nodes, edges_str, node_features_str)
@@ -100,19 +92,18 @@ def predict_graph(num_nodes, edges_str, node_features_str):
 # ---------- INTERFAZ GRADIO ----------
 
 description = """
-Este clasificador usa un modelo GCN entrenado sobre el dataset **MUTAG** para predecir si una molécula (representada como grafo) es mutagénica o no.
-
-🔹 Puedes definir tu propio grafo ingresando el número de nodos, las aristas y las características de cada nodo.
+Clasificador molecular basado en **Redes Neuronales en Grafo (GNN)** entrenado sobre el dataset **MUTAG**.
 
-✅ Cada nodo debe tener **7 características** (como en MUTAG).  
-🔗 Las aristas deben estar en formato Python: `[(0, 1), (1, 2)]`  
-📊 Las características deben ser una lista de listas: `[[1,0,0,1,0,1,0], [0,1,1,0,1,0,1], ...]`
+✏️ Puedes modificar el grafo:  
+- Número de nodos  
+- Aristas (formato: `[(0,1),(1,2)]`)  
+- Características de cada nodo (7 valores binarios por nodo)  
 """
 
 inputs = [
     gr.Number(label="Número de nodos", value=3, precision=0),
-    gr.Textbox(label="Aristas [(0,1), (1,2)]", lines=2, value="[(0,1),(1,2)]"),
-    gr.Textbox(label="Características por nodo (listas de 7)", lines=4, value="[[1,0,0,1,0,1,0], [0,1,1,0,1,0,1], [1,1,0,0,1,0,1]]")
+    gr.Textbox(label="Aristas", value="[(0,1),(1,2)]"),
+    gr.Textbox(label="Características por nodo", value="[[1,0,0,1,0,1,0],[0,1,1,0,1,0,1],[1,1,0,0,1,0,1]]")
 ]
 
 outputs = [
@@ -124,7 +115,7 @@ demo = gr.Interface(
     fn=predict_graph,
     inputs=inputs,
     outputs=outputs,
-    title="🔬 Clasificador Molecular con GNN (GCN)",
+    title="🧪 Clasificador Molecular con GCN",
     description=description,
     examples=[
         [3, "[(0,1),(1,2)]", "[[1,0,0,1,0,1,0],[0,1,1,0,1,0,1],[1,1,0,0,1,0,1]]"],
@@ -132,4 +123,4 @@ demo = gr.Interface(
     ]
 )
 
-demo.launch()
+demo.launch(show_error=True, cache_examples=False)