Spaces:
Sleeping
Sleeping
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"""<!DOCTYPE html><html> {output} </html>""" # do not use ' in this input | |
return f"""<iframe style="width: 100%; height:420px" name="result" allow="midi; geolocation; microphone; camera; | |
display-capture; encrypted-media;" sandbox="allow-modals allow-forms | |
allow-scripts allow-same-origin allow-popups | |
allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" | |
allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>""" | |
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() | |