v2.0.1

AstraBert · May 28, 2024 · b0a99f6 · b0a99f6
1 parent cb4762d
commit b0a99f6
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 147 deletions.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -8,7 +8,7 @@ WORKDIR /app
 ADD . /app
 
 #Upgrade gradio
-RUN pip install autotrain-advanced
+RUN pip install gradio_molecule3d
 
 # Expose the port that the application will run on
 EXPOSE 8760

diff --git a/docker/autotrain_interface.py b/docker/autotrain_interface.py
@@ -25,11 +25,7 @@ def build_command(hf_usr, hf_token, configpath):
             lines=3,
             value=f"your-powerful-token",
         ),
-        gr.Textbox(
-            label="Yaml configuration file",
-            info="Path to the yaml configuration file containing the information to use autotrain",
-            lines=3,
-            value="/path/to/config.yaml",
+        gr.File(label="Yaml configuration file path"
         )
     ],
     title="everything-ai-autotrain",

diff --git a/docker/protein_folding_with_esm.py b/docker/protein_folding_with_esm.py
@@ -1,8 +1,85 @@
 from transformers import AutoTokenizer, EsmForProteinFolding
 from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
 from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
-from proteins_viz import *
 import gradio as gr
+from gradio_molecule3d import Molecule3D
+
+reps =    [
+    {
+      "model": 0,
+      "chain": "",
+      "resname": "",
+      "style": "stick",
+      "color": "whiteCarbon",
+      "residue_range": "",
+      "around": 0,
+      "byres": False,
+      "visible": False
+    }
+]
+
+def read_mol(molpath):
+    with open(molpath, "r") as fp:
+        lines = fp.readlines()
+    mol = ""
+    for l in lines:
+        mol += l
+    return mol
+
+
+def molecule(input_pdb):
+
+    mol = read_mol(input_pdb)
+
+    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: 600px;
+    position: relative;
+    }
+    .mol-container select{
+        background-image:None;
+    }
+    </style>
+     <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+    <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>  
+    <div id="container" class="mol-container"></div>
+  
+            <script>
+               let pdb = `"""
+        + mol
+        + """`  
+      
+             $(document).ready(function () {
+                let element = $("#container");
+                let config = { backgroundColor: "white" };
+                let viewer = $3Dmol.createViewer(element, config);
+                viewer.addModel(pdb, "pdb");
+                viewer.getModel(0).setStyle({}, { cartoon: { colorscheme:"whiteCarbon" } });
+                viewer.zoomTo();
+                viewer.render();
+                viewer.zoom(0.8, 2000);
+              })
+        </script>
+        </body></html>"""
+    )
+
+    return f"""<iframe style="width: 100%; height: 600px" 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 convert_outputs_to_pdb(outputs):
     final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
@@ -47,14 +124,19 @@ def fold_protein(test_protein):
     pdb = convert_outputs_to_pdb(output)
     with open("output_structure.pdb", "w") as f:
         f.write("".join(pdb))
-    image = take_care("output_structure.pdb")
-    return image
+    html = molecule("output_structure.pdb")
+    return html, "output_structure.pdb"
 
 iface = gr.Interface(
     title="everything-ai-proteinfold",
     fn=fold_protein,
-    inputs="text",
-    outputs="image", 
+    inputs=gr.Textbox(
+            label="Protein Sequence",
+            info="Find sequences examples below, and complete examples with images at: https://github.com/AstraBert/proteinviz/tree/main/examples.md; if you input a sequence, you're gonna get the static image and the 3D model to explore and play with",
+            lines=5,
+            value=f"Paste or write amino-acidic sequence here",
+        ),
+    outputs=[gr.HTML(label="Protein 3D model"), Molecule3D(label="Molecular 3D model", reps=reps)], 
 )
 
-iface.launch(server_name="0.0.0.0", share=False)
+iface.launch(server_name="0.0.0.0", share=False)
diff --git a/docker/proteins_viz.py b/docker/proteins_viz.py