Update app.py

app.py

@@ -1,30 +1,49 @@
 
+import os
+import json
+
 import gradio as gr
 from gradio_moleculeview import moleculeview
 
 
-
+import cellscape
 
 def predict(input_mol, view_str, chains):
-    return "svg"
+    # write view to file
+    with open("view_matrix", "w") as f:
+        f.write(view_str)
+
+    chain_str = ""
+    chain_dict = json.reads(chains)
+
+    # sort keys in dict and add colors to chain_str
+    for chain in sorted(chain_dict.keys()):
+        chain_str += f" {chain_dict[chain]}"
+
+    os.system(f"cellscape cartoon --pdb {input_mol.name} --outline chain --colors {chain_str}  --depth flat --back_outline --view view_matrix --save outline_all.svg")
+
+    #read content of file
+    with open("outline_all.svg", "r") as f:
+        return f.read(), "outline_all.svg"
 
 
 
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Molecule3D")
+    gr.Markdown("# PDB2Vector")
     inp = moleculeview(label="Molecule3D")
 
     view_str = gr.Textbox("viewMatrixResult", label="View Matrix", visible=False)
     chains = gr.Textbox("chainsResult", label="Chains", visible=False) 
     
 
-    btn = gr.Button("Predict")
+    btn = gr.Button("Vectorize")
     html = gr.HTML("")
+    out_file = gr.File(label="Download SVG")
     btn.click(None, [], [view_str, chains], js="() => [document.getElementById('viewMatrixResult').value, document.getElementById('chains').value]")  #
     # on change of chains trigger, rendering
-    chains.change(predict, [inp, view_str, chains], html)
+    chains.change(predict, [inp, view_str, chains], [html, out_file])
     
 
 if __name__ == "__main__":