app.py

import os
import json

import gradio as gr
from gradio_moleculeview import moleculeview


import cellscape

def predict(input_mol, style, view_str, chains):
    # write view to file
    with open("view_matrix", "w") as f:
        f.write(json.loads(view_str))

    chain_str = ""
    chain_dict = json.loads(chains)

    # sort keys in dict and add colors to chain_str
    for chain in sorted(chain_dict.keys()):
        chain_str += f" '{chain_dict[chain]}'"
    if style == "Goodsell3D":
        os.system(f"cellscape cartoon --pdb {input_mol.name} --outline residue --depth_shading --depth_lines --colors {chain_str}  --depth flat --back_outline --view view_matrix --save outline_all.svg")
    elif style == "Contour":
        os.system(f"cellscape cartoon --pdb {input_mol.name} --outline residue  --depth_contour_interval 10 --colors {chain_str}  --depth contours --back_outline --view view_matrix --save outline_all.svg")
    else:
        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
    print(os.stat("outline_all.svg").st_size / (1024 * 1024))
    os.system("inkscape outline_all.svg --actions='select-all;path-simplify;export-plain-svg' --export-filename pdb_opt.svg")
    print(os.stat("outline_all.svg").st_size / (1024 * 1024))

    with open("pdb_opt.svg", "r") as f:
        return f.read(), "pdb_opt.svg"





with gr.Blocks() as demo:
    gr.Markdown("# PDB2Vector")
    style = gr.Radio(value="Flat", choices=["Flat", "Contour", "Goodsell3D"], label="Style")
    inp = moleculeview(label="Molecule3D")

    view_str = gr.Textbox("viewMatrixResult", label="View Matrix", visible=False)
    chains = gr.Textbox("chainsResult", label="Chains", visible=False) 
    hidden_style = gr.Textbox(visible=False)

    btn = gr.Button("1. Save state")
    btn2 = gr.Button("2. Vectorize")
    html = gr.HTML("")
    out_file = gr.File(label="Download SVG")
    btn.click(None, style, [view_str, chains, hidden_style], js="(style) => [document.getElementById('viewMatrixResult').value, document.getElementById('chains').value, style]")  #
    btn2.click(predict, [inp, style, view_str, chains], [html, out_file])
    # on change of chains trigger, rendering
    # view_str.change(predict, [inp, style, view_str, chains], [html, out_file])
    # chains.change(predict, [inp, style, view_str, chains], [html, out_file])
    # hidden_style.change(predict, [inp, style, view_str, chains], [html, out_file])
    

if __name__ == "__main__":
    demo.launch()