Update app.py

app.py

@@ -2,10 +2,15 @@ import torch
 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
+import io
+import base64
+
+# ---------- MODELO GCN ----------
 
-# --- Mismo modelo que en Colab ---
 class GCN(torch.nn.Module):
     def __init__(self, hidden_channels=64):
         super().__init__()
@@ -21,20 +26,63 @@ class GCN(torch.nn.Module):
         x = global_mean_pool(x, batch)
         return self.lin(x)
 
-# --- Carga de modelo entrenado ---
+# ---------- CARGA DEL MODELO ENTRENADO ----------
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = GCN().to(device)
 model.load_state_dict(torch.load("model_gcn.pth", map_location=device))
 model.eval()
 
-# --- Función de predicción sobre un grafo ejemplo ---
-def demo_predict():
+# ---------- FUNCIÓN PARA PARSEAR INPUT ----------
+
+def parse_input(num_nodes, edges_str, node_features_str):
     G = nx.Graph()
-    G.add_edges_from([(0, 1), (1, 2)])
-    nx.set_node_attributes(G, {i: [1, 0, 0, 1, 0, 1, 0] for i in G.nodes}, "x")  # vector de 7 dimensiones
+
+    try:
+        # Añadimos nodos
+        for i in range(num_nodes):
+            G.add_node(i)
+
+        # Parseamos aristas
+        edges = eval(edges_str)  # formato esperado: [(0,1), (1,2)]
+        G.add_edges_from(edges)
+
+        # Parseamos características
+        node_features = eval(node_features_str)  # formato: [[1,0,1,0,1,0,1], [...]]
+        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.")
+
+        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 ----------
+
+def draw_graph(G, pred_label):
+    pos = nx.spring_layout(G)
+    node_colors = 'lightgreen' if pred_label == 1 else 'lightcoral'
+    plt.figure(figsize=(4, 4))
+    nx.draw(G, pos, with_labels=True, node_color=node_colors, edge_color='gray', node_size=800)
+    plt.title("Grafo de entrada")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    img_base64 = base64.b64encode(buf.read()).decode('utf-8')
+    plt.close()
+
+    return f"data:image/png;base64,{img_base64}"
+
+# ---------- FUNCIÓN DE PREDICCIÓN ----------
+
+def predict_graph(num_nodes, edges_str, node_features_str):
+    G = parse_input(num_nodes, edges_str, node_features_str)
 
     data = from_networkx(G)
-    data.x = torch.tensor(list(nx.get_node_attributes(G, "x").values()), dtype=torch.float)
+    data.x = torch.tensor([v for v in nx.get_node_attributes(G, "x").values()], dtype=torch.float)
     data.edge_index = data.edge_index
     data.batch = torch.tensor([0] * data.num_nodes)
 
@@ -43,15 +91,43 @@ def demo_predict():
         out = model(data.x, data.edge_index, data.batch)
         pred = out.argmax(dim=1).item()
 
-    return "Mutagénico ✅" if pred == 1 else "No mutagénico ❌"
+    label_text = "Mutagénico ✅" if pred == 1 else "No mutagénico ❌"
+    img = draw_graph(G, pred)
+    return label_text, img
+
+# ---------- 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.
+
+✅ 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], ...]`
+"""
+
+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]]")
+]
+
+outputs = [
+    gr.Text(label="Predicción"),
+    gr.Image(label="Visualización del grafo")
+]
 
-# --- Interfaz Gradio ---
 demo = gr.Interface(
-    fn=demo_predict,
-    inputs=[],
-    outputs="text",
-    title="Clasificador de Moléculas con GCN",
-    description="Este demo usa una red neuronal en grafo entrenada sobre MUTAG para clasificar moléculas como mutagénicas o no."
+    fn=predict_graph,
+    inputs=inputs,
+    outputs=outputs,
+    title="🔬 Clasificador Molecular con GNN (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]]"],
+        [4, "[(0,1),(1,2),(2,3)]", "[[1,1,0,1,0,0,1],[0,0,1,1,1,0,0],[1,0,1,0,1,1,0],[0,1,0,1,1,0,1]]"]
+    ]
 )
 
 demo.launch()