Spaces:

silterra
/

DiffDock-Pocket-Web

Sleeping

App Files Files Community

silterra commited on Nov 14, 2023

Commit

b1dc6ad

•

1 Parent(s): 5c0a4b6

Display ligand (SDF) file as well as protein.

Browse files

Files changed (3) hide show

main.py +3 -3
mol_viewer.py +18 -11
run_utils.py +15 -7

main.py CHANGED Viewed

@@ -16,13 +16,13 @@ def run_wrapper(protein_file, ligand_file, other_args_file, *args, **kwargs) ->
     if other_args_file is not None:
         kwargs["other_arg_file"] = other_args_file.name
-    output_file, output_pdb_text = run_utils.run_cli_command(
         protein_file.name, ligand_file.name, *args, **kwargs
     )
     output_viz = "No visualisation created"
-    if output_pdb_text:
-        output_viz = mol_viewer.gen_3dmol_vis(output_pdb_text)
     message = f"Calculation completed at {datetime.datetime.now()}"
     return message, output_file, output_viz

     if other_args_file is not None:
         kwargs["other_arg_file"] = other_args_file.name
+    output_file, pdb_text, sdf_text = run_utils.run_cli_command(
         protein_file.name, ligand_file.name, *args, **kwargs
     )
     output_viz = "No visualisation created"
+    if pdb_text:
+        output_viz = mol_viewer.gen_3dmol_vis(pdb_text, sdf_text)
     message = f"Calculation completed at {datetime.datetime.now()}"
     return message, output_file, output_viz

mol_viewer.py CHANGED Viewed

@@ -9,9 +9,7 @@ https://huggingface.co/spaces/simonduerr/3dmol.js/blob/main/app.py
 """
-def gen_3dmol_vis(pdb_text: str):
-    mol = pdb_text
     x = (
             """<!DOCTYPE html>
             <html>
@@ -38,26 +36,35 @@ def gen_3dmol_vis(pdb_text: str):
         <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 gen_3dmol_vis(pdb_text: str, sdf_text: str):
     x = (
             """<!DOCTYPE html>
             <html>
         <div id="container" class="mol-container"></div>
                 <script>
+                   let pdb = `""" + pdb_text + """`
+                   let sdf = `""" + sdf_text + """`
              $(document).ready(function () {
                 let element = $("#container");
                 let config = { backgroundColor: "white" };
                 let viewer = $3Dmol.createViewer(element, config);
+                viewer.addModel(pdb, "pdb", { format: "pdb" });
+                pdb_model = viewer.getModel(0);
+                // Cartoon view for protein
+                // First argument is a selector
+                pdb_model.setStyle({}, {cartoon: {color: "spectrum", opacity: 0.8}});
+                viewer.addModel(sdf, "sdf", {format: "sdf"});
+                sdf_model = viewer.getModel(1);
+                // Stick view for SDF
+                sdf_model.setStyle({}, {stick: {color: "red", radius: 0.2, opacity: 0.8}});
                 viewer.zoomTo();
                 viewer.render();
+                // viewer.zoom(0.8, 2000);
               })
         </script>
         </body></html>"""
     )
+    return f"""<iframe style="width: 100%; height: 600px" name="result" 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>"""

run_utils.py CHANGED Viewed

@@ -124,15 +124,23 @@ def run_cli_command(
         # If there's a file for viewing, load it and view it.
         sub_dirs = [os.path.join(temp_dir_path, x) for x in os.listdir(temp_dir_path)]
         sub_dirs = list(filter(lambda x: os.path.isdir(x), sub_dirs))
-        display_path = display_pdb_text = None
         if len(sub_dirs) == 1:
-            display_path = os.path.join(sub_dirs[0], "rank1_reverseprocess_protein.pdb")
-        if display_path and os.path.exists(display_path):
-            display_pdb_text = read_file_lines(display_path)
-        logging.debug(f"Display path: {display_path}")
-        # logging.debug(f"Display pdb text:\n{display_pdb_text}")
         # Zip the output directory
         # Generate a unique filename using a timestamp and a UUID
@@ -144,7 +152,7 @@ def run_cli_command(
         logging.debug(f"Directory '{temp_dir}' zipped to '{full_zip_path}'")
-    return full_zip_path, display_pdb_text
 if __name__ == "__main__":

         # If there's a file for viewing, load it and view it.
         sub_dirs = [os.path.join(temp_dir_path, x) for x in os.listdir(temp_dir_path)]
         sub_dirs = list(filter(lambda x: os.path.isdir(x), sub_dirs))
+        pdb_path = pdb_text = sdf_path = sdf_text = None
         if len(sub_dirs) == 1:
+            sub_dir = sub_dirs[0]
+            pdb_path = os.path.join(sub_dir, "rank1_reverseprocess_protein.pdb")
+            sdf_path = os.path.join(sub_dir, "rank1.sdf")
+        if skip_running:
+            # Test/debugging only
+            example_dir = os.path.join(os.environ["HOME"], "Projects", "DiffDock-Pocket", "example_data", "example_outputs")
+            pdb_path = os.path.join(example_dir, "rank1_reverseprocess_protein.pdb")
+            sdf_path = os.path.join(example_dir, "rank1.sdf")
+        if pdb_path and os.path.exists(pdb_path):
+            pdb_text = read_file_lines(pdb_path)
+            sdf_text = read_file_lines(sdf_path)
+        logging.debug(f"Display path: {pdb_path}")
         # Zip the output directory
         # Generate a unique filename using a timestamp and a UUID
         logging.debug(f"Directory '{temp_dir}' zipped to '{full_zip_path}'")
+    return full_zip_path, pdb_text, sdf_text
 if __name__ == "__main__":