import gradio as gr import py3Dmol import torch import pydssp 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 run(pdb_id, pdb_file, helix, sheet, loop): path_to_pdb = get_pdb(pdb_code=pdb_id, 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, "-": loop, "E": sheet} 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: pdb_id = gr.Textbox(label="pdb code") pdb_file = gr.File(label="pdb file") 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() gr.Examples( [["1QYS", "#ff0000", "#00ff00", "#0000ff"]], inputs=[pdb_id, helix, sheet, loop], outputs=[html], fn=run, ) btn.click(fn=run, inputs=[pdb_id, pdb_file, helix, sheet, loop], outputs=[html]) demo.launch()