Spaces:
Running
Running
| import gradio as gr | |
| import urllib | |
| import re | |
| import sys | |
| import warnings | |
| import torch | |
| import torch.nn as nn | |
| import ipywidgets as widgets | |
| from ipywidgets import interact, fixed | |
| from utils.helpers import * | |
| from utils.voxelization import processStructures | |
| from utils.model import Model | |
| import numpy as np | |
| import os | |
| import moleculekit | |
| print(moleculekit.__version__) | |
| def update(inp, file, mode): | |
| try: | |
| pdb_file = file.name | |
| except: | |
| print("using pdbfile") | |
| try: | |
| pdb_file = inp | |
| if ( | |
| re.match( | |
| "[OPQ][0-9][A-Z0-9]{3}[0-9]|[A-NR-Z][0-9]([A-Z][A-Z0-9]{2}[0-9]){1,2}", | |
| pdb_file, | |
| ).group() | |
| == pdb_file | |
| ): | |
| urllib.request.urlretrieve( | |
| f"https://alphafold.ebi.ac.uk/files/AF-{pdb_file}-F1-model_v2.pdb", | |
| f"files/{pdb_file}.pdb", | |
| ) | |
| except AttributeError: | |
| if len(inp) == 4: | |
| pdb_file = inp | |
| urllib.request.urlretrieve( | |
| f"http://files.rcsb.org/download/{pdb_file.lower()}.pdb1", | |
| f"files/{pdb_file}.pdb", | |
| ) | |
| else: | |
| return "pdb code must be 4 letters or Uniprot code does not match", "" | |
| if mode == "All residues": | |
| ids = get_all_protein_resids( | |
| f"files/{pdb_file}.pdb", | |
| ) | |
| else: | |
| ids = get_all_metalbinding_resids(f"files/{pdb_file}.pdb") | |
| voxels, prot_centers, prot_N, prots = processStructures(pdb_file, ids) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| voxels.to(device) | |
| print(voxels.shape) | |
| model = Model() | |
| model.to(device) | |
| model.load_state_dict(torch.load("weights/metal_0.5A_v3_d0.2_16Abox.pth", map_location=torch.device('cpu'))) | |
| model.eval() | |
| with warnings.catch_warnings(): | |
| warnings.filterwarnings("ignore") | |
| output = model(voxels) | |
| print(output.shape) | |
| prot_v = np.vstack(prot_centers) | |
| output_v = output.flatten().cpu().detach().numpy() | |
| bb = get_bb(prot_v) | |
| gridres = 0.5 | |
| grid, box_N = create_grid_fromBB(bb, voxelSize=gridres) | |
| probability_values = get_probability_mean(grid, prot_v, output_v) | |
| print(probability_values.shape) | |
| write_cubefile( | |
| bb, | |
| probability_values, | |
| box_N, | |
| outname=f"output/metal_{pdb_file}.cube", | |
| gridres=gridres, | |
| ) | |
| message = find_unique_sites( | |
| probability_values, | |
| grid, | |
| writeprobes=True, | |
| probefile=f"output/probes_{pdb_file}.pdb", | |
| threshold=7, | |
| p=0.15, | |
| ) | |
| return message, molecule( | |
| f"files/{pdb_file}.pdb", | |
| f"output/probes_{pdb_file}.pdb", | |
| f"output/metal_{pdb_file}.cube", | |
| ) | |
| def test(): | |
| x = """<!DOCTYPE html> | |
| <html> | |
| <head> | |
| <meta http-equiv="content-type" content="text/html; charset=UTF-8" /> | |
| </head> | |
| <body> | |
| <script src="https://3Dmol.org/build/3Dmol-min.js" async></script> <div style="height: 400px; width: 400px; position: relative;" class="viewer_3Dmoljs" data-pdb="2POR" data-backgroundcolor="0xffffff" data-style="stick" ></div> | |
| </body></html>""" | |
| return f"""<iframe style="width: 100%; height: 480px" 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>""" | |
| def read_mol(molpath): | |
| with open(molpath, "r") as fp: | |
| lines = fp.readlines() | |
| mol = "" | |
| for l in lines: | |
| mol += l | |
| return mol | |
| def molecule(pdb, probes, cube): | |
| mol = read_mol(pdb) | |
| probes = read_mol(probes) | |
| cubefile = read_mol(cube) | |
| x = ( | |
| """<!DOCTYPE html> | |
| <html> | |
| <head> | |
| <meta http-equiv="content-type" content="text/html; charset=UTF-8" /> | |
| <style> | |
| body{ | |
| font-family:sans-serif | |
| } | |
| .mol-container { | |
| width: 100%; | |
| height: 400px; | |
| position: relative; | |
| } | |
| .slider{ | |
| width:80%; | |
| margin:0 auto | |
| } | |
| .slidercontainer{ | |
| display:flex; | |
| } | |
| .slidercontainer > * + * { | |
| margin-left: 0.5rem; | |
| } | |
| #isovalue{ | |
| text-align:right} | |
| </style> | |
| <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script> | |
| <script src="https://cdnjs.cloudflare.com/ajax/libs/rangeslider.js/2.3.3/rangeslider.min.js" integrity="sha512-BUlWdwDeJo24GIubM+z40xcj/pjw7RuULBkxOTc+0L9BaGwZPwiwtbiSVzv31qR7TWx7bs6OPTE5IyfLOorboQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script> | |
| </head> | |
| <body> | |
| <div class="slidercontainer"> | |
| <span>Isovalue </span> | |
| <span id="isovalue">0.5</span> | |
| <input class="slider" type="range" id="rangeslider" min="0" max="1" step="0.025" value=0.5> | |
| </div> | |
| <div id="container" class="mol-container"></div> | |
| <script> | |
| let viewer = null; | |
| let voldata = null; | |
| $(document).ready(function () { | |
| let element = $("#container"); | |
| let config = { backgroundColor: "white" }; | |
| viewer = $3Dmol.createViewer( element, config ); | |
| viewer.ui.initiateUI(); | |
| let data = `""" | |
| + mol | |
| + """` | |
| viewer.addModel( data, "pdb" ); | |
| let cubefile = `""" | |
| + cubefile | |
| + """` | |
| voldata = new $3Dmol.VolumeData(cubefile, "cube"); | |
| viewer.addIsosurface(voldata, { isoval: 0.7 , color: "blue", alpha: 0.85, smoothness: 1 }); | |
| viewer.getModel(0).setStyle({}, {cartoon: {color: "grayCarbon"}}); | |
| let probes =`""" | |
| + probes | |
| + """` | |
| viewer.addModel(probes, "pdb"); | |
| viewer.getModel(1).setStyle({ "resn": "ZN" }, { "sphere": { }}); | |
| viewer.getModel(1).setHoverable({}, true, | |
| function (atom, viewer, event, container) { | |
| if (!atom.label) { | |
| atom.label = viewer.addLabel("ZN p=" + atom.pdbline.substring(55, 60), { position: atom, backgroundColor: "mintcream", fontColor: "black" }); | |
| } | |
| }, | |
| function (atom, viewer) { | |
| if (atom.label) { | |
| viewer.removeLabel(atom.label); | |
| delete atom.label; | |
| } | |
| } | |
| ); | |
| viewer.zoomTo(); | |
| viewer.render(); | |
| viewer.zoom(0.8, 2000); | |
| }); | |
| </script> | |
| <script> | |
| $("#rangeslider").rangeslider().on("change", function (el) { | |
| isoval = parseFloat(el.target.value); | |
| $("#isovalue").text(el.target.value) | |
| viewer.addIsosurface(voldata, { isoval: parseFloat(el.target.value), color: "blue", alpha: 0.85, smoothness: 1 }); | |
| viewer.render(); | |
| }); | |
| </script> | |
| </body></html>""" | |
| ) | |
| return f"""<iframe style="width: 100%; height: 480px" 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>""" | |
| metal3d = gr.Blocks() | |
| with metal3d: | |
| gr.Markdown("# Metal3D") | |
| with gr.Group(): | |
| inp = gr.Textbox( | |
| placeholder="PDB Code or Uniprot identifier or upload file below", label="Input molecule" | |
| ) | |
| file = gr.File(file_count="single", type="file") | |
| with gr.Row(): | |
| mode = gr.Radio( | |
| ["All metalbinding residues (ASP, CYS, GLU, HIS)", "All residues"], | |
| label="Residues to use for prediction", | |
| ) | |
| custom_resids = gr.Textbox(placeholder="Comma separated list of residues", label="Custom residues") | |
| btn = gr.Button("Run") | |
| gr.Markdown( | |
| """ <small>Inference using CPU-only, can be quite slow for more than 20 residues. Use Colab notebook for GPU acceleration</small> | |
| """ | |
| ) | |
| gr.Markdown("# Output") | |
| out = gr.Textbox(label="status") | |
| mol = gr.HTML() | |
| btn.click(fn=update, inputs=[inp, file, mode], outputs=[out, mol]) | |
| metal3d.launch() | |