Upload app.py

app.py

@@ -1,11 +1,18 @@
 import os
 os.system("pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu")
 os.system("pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.12.0+cpu.html")
+os.system("apt-get install -y graphviz-dev")
+os.system("pip install pygraphviz")
+
 import gradio as gr
 from glycowork.ml.processing import dataset_to_dataloader
 import numpy as np
 import torch
 import torch.nn as nn
+from networkx.drawing.nx_agraph import write_dot
+import pygraphviz as pgv
+from glycowork.motif.graph import glycan_to_nxGraph
+import networkx as nx
 
 class EnsembleModel(nn.Module):
     def __init__(self, models):
@@ -33,6 +40,17 @@ model2 = torch.load("model2.pt", map_location=torch.device('cpu'))
 model3 = torch.load("model3.pt", map_location=torch.device('cpu'))
 
 def fn(glycan, model):
+    # Draw graph
+    graph = glycan_to_nxGraph(glycan)
+    node_labels = nx.get_node_attributes(graph, 'string_labels')
+    labels = {i:node_labels[i] for i in range(len(graph.nodes))}
+    graph = nx.relabel_nodes(graph, labels)
+    write_dot(graph, "graph.dot")
+    graph=pgv.AGraph("graph.dot")  
+    graph.layout(prog='dot')
+    graph.draw("graph.png")
+    
+    # Perform inference
     if model == "No data augmentation":
       model_pred = model1
       model_pred.eval()
@@ -62,13 +80,13 @@ def fn(glycan, model):
     pred = np.exp(pred)/sum(np.exp(pred)) # Softmax 
     pred = [float(x) for x in pred]
     pred = {class_list[i]:pred[i] for i in range(15)}
-    return pred
+    return pred, "graph.png"
 
 
 demo = gr.Interface(
     fn=fn,
     inputs=[gr.Textbox(label="Glycan sequence"), gr.Radio(label="Model",choices=["No data augmentation", "Random node deletion", "Ensemble"])],
-    outputs=[gr.Label(num_top_classes=15, label="Prediction")],
+    outputs=[gr.Label(num_top_classes=15, label="Prediction"), gr.Image(label="Graph visualization")],
     allow_flagging=False,
     title="SweetNet demo",
     examples=[["GlcOSN(a1-4)GlcA(b1-4)GlcOSN(a1-4)GlcAOS(b1-4)GlcOSN(a1-4)GlcOSN", "No data augmentation"],