container

Dockerfile

@@ -51,16 +51,18 @@ ENV PYTHONPATH="$AMBERHOME/lib/python3.10/site-packages"
 #COPY --chown=user . $HOME/app
 
 
-#RUN adduser -u 5678 --disabled-password --gecos "" appuser && chown -R appuser .
-#USER appuser
-#CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
+RUN adduser -u 5678 --disabled-password --gecos "" appuser && chown -R appuser .
 
-WORKDIR $HOME/app
+USER appuser
 
-# Copy the current directory contents into the container at $HOME/app setting the owner to the user
-COPY --chown=user . $HOME/app
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
 
-CMD ["python", "app.py"]
+#WORKDIR $HOME/app
+#
+## Copy the current directory contents into the container at $HOME/app setting the owner to the user
+#COPY --chown=user . $HOME/app
+#
+#CMD ["python", "app.py"]
 
 
 

main.py

@@ -1,30 +1,150 @@
 import gradio as gr
+
+import py3Dmol
+
+from Bio.PDB import *
+
+import numpy as np
+from Bio.PDB import PDBParser
+import pandas as pd
 import torch
-import requests
-from torchvision import transforms
+import os
+from MDmodel import GNN_MD
+import h5py
+from transformMD import GNNTransformMD
+
+# JavaScript functions
+resid_hover = """function(atom,viewer) {{
+    if(!atom.label) {{
+        atom.label = viewer.addLabel('{0}:'+atom.atom+atom.serial,
+            {{position: atom, backgroundColor: 'mintcream', fontColor:'black'}});
+    }}
+}}"""
+hover_func = """
+function(atom,viewer) {
+    if(!atom.label) {
+        atom.label = viewer.addLabel(atom.interaction,
+            {position: atom, backgroundColor: 'black', fontColor:'white'});
+    }
+}"""
+unhover_func = """
+function(atom,viewer) {
+    if(atom.label) {
+        viewer.removeLabel(atom.label);
+        delete atom.label;
+    }
+}"""
+atom_mapping = {0:'H', 1:'C', 2:'N', 3:'O', 4:'F', 5:'P', 6:'S', 7:'CL', 8:'BR', 9:'I', 10: 'UNK'}
+
+model = GNN_MD(11, 64)
+state_dict = torch.load(
+    "best_weights_rep0.pt",
+    map_location=torch.device("cpu"),
+)["model_state_dict"]
+model.load_state_dict(state_dict)
+model = model.to('cpu')
+model.eval()
+
+
+
+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):
+    #path_to_pdb = get_pdb(pdb_code=pdb_code, filepath=pdb_file)
+
+    #pdb = open(path_to_pdb, "r").read()
+    # switch to misato env if not running from container
+    mdh5_file = "inference_for_md.hdf5"
+    md_H5File = h5py.File(mdh5_file)
+
+    column_names = ["x", "y", "z", "element"]
+    atoms_protein = pd.DataFrame(columns = column_names)
+    cutoff = md_H5File["11GS"]["molecules_begin_atom_index"][:][-1] # cutoff defines protein atoms
+
+    atoms_protein["x"] = md_H5File["11GS"]["atoms_coordinates_ref"][:][:cutoff, 0]
+    atoms_protein["y"] = md_H5File["11GS"]["atoms_coordinates_ref"][:][:cutoff, 1]
+    atoms_protein["z"] = md_H5File["11GS"]["atoms_coordinates_ref"][:][:cutoff, 2]
+
+    atoms_protein["element"] = md_H5File["11GS"]["atoms_element"][:][:cutoff]  
+
+    item = {}
+    item["scores"] = 0
+    item["id"] = "11GS"
+    item["atoms_protein"] = atoms_protein
+
+    transform = GNNTransformMD()
+    data_item = transform(item)
+    adaptability = model(data_item)
+    adaptability = adaptability.detach().numpy()
+    
+    data = []
+
+
+    for i in range(adaptability.shape[0]):
+        data.append([i, atom_mapping(atoms_protein.iloc[i, atoms_protein.columns.get_loc("element")] - 1), atoms_protein.iloc[i, atoms_protein.columns.get_loc("x")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("y")],atoms_protein.iloc[i, atoms_protein.columns.get_loc("z")],adaptability[i]])
+
+    topN = 100
+    topN_ind = np.argsort(adaptability)[::-1][:topN]    
+
+    pdb = open(pdb_file.name, "r").read()
+
+    view = py3Dmol.view(width=600, height=400)
+    view.setBackgroundColor('black')
+    view.addModel(pdb, "pdb")
+    view.setStyle({'stick': {'colorscheme': {'prop': 'resi', 'C': 'turquoise'}}})
+   
+    for i in range(topN):
+        view.addSphere({'center':{'x':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("x")], 'y':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("y")],'z':atoms_protein.iloc[topN_ind[i], atoms_protein.columns.get_loc("z")]},'radius':adaptability[topN_ind[i]]/1.5,'color':'orange','alpha':0.75})    
 
-model = torch.hub.load("pytorch/vision:v0.6.0", "resnet18", pretrained=True).eval()
-response = requests.get("https://git.io/JJkYN")
-labels = response.text.split("\n")
+    view.zoomTo()
 
+    output = view._make_html().replace("'", '"')
 
-def predict(inp):
-    inp = transforms.ToTensor()(inp).unsqueeze(0)
-    with torch.no_grad():
-        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
-        confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
-    return confidences
+    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>""", pd.DataFrame(data, columns=['index','element','x','y','z','Adaptability'])
 
 
-def run():
-    demo = gr.Interface(
-        fn=predict,
-        inputs=gr.inputs.Image(type="pil"),
-        outputs=gr.outputs.Label(num_top_classes=3),
-    )
+callback = gr.CSVLogger()
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+with gr.Blocks() as demo:
+    gr.Markdown("# Protein Adaptability Prediction")
+    
+    #text_input = gr.Textbox()
+    #text_output = gr.Textbox()
+    #text_button = gr.Button("Flip")
+    inp = gr.Textbox(placeholder="PDB Code or upload file below", label="Input structure")
+    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")
+    single_btn = gr.Button(label="Run")
+    with gr.Row():
+        html = gr.HTML()
+    with gr.Row():
+        dataframe = gr.Dataframe()
+            
+    single_btn.click(fn=predict, inputs=[inp, pdb_file], outputs=[html, dataframe])
 
 
-if __name__ == "__main__":
-    run()
+demo.launch(debug=True)

old_main.py

@@ -0,0 +1,30 @@
+import gradio as gr
+import torch
+import requests
+from torchvision import transforms
+
+model = torch.hub.load("pytorch/vision:v0.6.0", "resnet18", pretrained=True).eval()
+response = requests.get("https://git.io/JJkYN")
+labels = response.text.split("\n")
+
+
+def predict(inp):
+    inp = transforms.ToTensor()(inp).unsqueeze(0)
+    with torch.no_grad():
+        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
+        confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
+    return confidences
+
+
+def run():
+    demo = gr.Interface(
+        fn=predict,
+        inputs=gr.inputs.Image(type="pil"),
+        outputs=gr.outputs.Label(num_top_classes=3),
+    )
+
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+
+
+if __name__ == "__main__":
+    run()