import os import json import gradio as gr from gradio_moleculeview import moleculeview import cellscape def html_output(input_file): with open(input_file, "r") as f: svg = f.read().replace(" """ + svg + """ """ ) x = x.replace('"', """) x = x.replace("'", """) return f"""""" def predict(input_mol, style, contour_level, 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 --color_by chain --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 chain --color_by chain --depth_contour_interval {contour_level} --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)) return html_output("outline_all.svg"), "pdb_opt.svg" def show_contour_level(style): if style=="Contour": return gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=True) else: return gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=False) with gr.Blocks() as demo: gr.Markdown("# PDB2Vector") style = gr.Radio(value="Flat", choices=["Flat", "Contour", "Goodsell3D"], label="Style") contour_level = gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=False) style.change(show_contour_level, style, contour_level) 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) timestamp = gr.Textbox(visible=False) btn = gr.Button("Vectorize") html = gr.HTML("") out_file = gr.File(label="Download SVG") btn.click(None, style, [view_str, chains, hidden_style, timestamp], js="(style) => [document.getElementById('viewMatrixResult').value, document.getElementById('chains').value, style, Date.now()]") # timestamp.change(predict, [inp, style, contour_level, view_str, chains], [html, out_file]) # on change of chains trigger, rendering if __name__ == "__main__": demo.launch()