import gradio as gr import torch import pydssp import py3Dmol import os def get_pdb(pdb_code="", filepath=""): try: return filepath.name except AttributeError as e: if pdb_code is None or pdb_code == "": return None else: os.system(f"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb") return f"{pdb_code}.pdb" def get_offset(pdb): pdb_multiline = pdb.split("\n") for line in pdb_multiline: if line.startswith("ATOM"): return int(line[22:27]) def predict(pdb_code, pdb_file, helix, sheet, loop): path_to_pdb = get_pdb(pdb_code=pdb_code, filepath=pdb_file) pdb = open(path_to_pdb, "r").read() offset = get_offset(pdb) coord = torch.tensor(pydssp.read_pdbtext(pdb)) secondary_struct = pydssp.assign(coord) view = py3Dmol.view(width=400, height=400) view.addModel(pdb, "pdb") colormap = {"H": helix, "E": sheet, "-": loop} colors = {i + offset: colormap[resi] for i, resi in enumerate(secondary_struct)} view.setStyle({"cartoon": {"colorscheme": {"prop": "resi", "map": colors}}}) view.zoomTo() output = view._make_html().replace("'", '"') x = f""" {output} """ # do not use ' in this input return f"""""" with gr.Blocks() as demo: gr.Markdown("# Secondary structure using pyDSSP") pdb_code = gr.Textbox(label="PDB Code") pdb_file = gr.File(label="PDB File Upload") with gr.Row(): helix = gr.ColorPicker(label="helix") sheet = gr.ColorPicker(label="sheet") loop = gr.ColorPicker(label="loop") btn = gr.Button(label="Run") html = gr.HTML() btn.click( fn=predict, inputs=[pdb_code, pdb_file, helix, sheet, loop], outputs=[html] ) gr.Examples( [["1QYS", "#ff0000", "#00ff00", "#0000ff"]], inputs=[pdb_code, helix, sheet, loop], fn=predict, ) demo.launch()