Complete overhaul of the UI.

* Use Blocks instead of Interface.

* Show protein in nice 3dmol viewer after calculation is finished.

main.py

@@ -1,54 +1,69 @@
-import logging
-import os.path
+import datetime
+from typing import Tuple, Optional
 
 import gradio as gr
 
+import mol_viewer
 import run_utils
 
 
-def run_wrapper(protein_file, ligand_file, other_args_file, *args, **kwargs) -> str:
+def run_wrapper(protein_file, ligand_file, other_args_file, *args, **kwargs) -> Tuple[str, Optional[str], str]:
+    if protein_file is None:
+        return "Protein file is missing! Must provide a protein file in PDB format", None, ""
+    if ligand_file is None:
+        return "Ligand file is missing! Must provide a ligand file in SDF format", None, ""
+
     if other_args_file is not None:
         kwargs["other_arg_file"] = other_args_file.name
-    return run_utils.run_cli_command(
+
+    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
+
 
 def run():
-    iface = gr.Interface(
-        fn=run_wrapper,
-        inputs=[
-            gr.File(label="Protein PDB", file_types=[".pdb"]),
-            gr.File(label="Ligand SDF", file_types=[".sdf"]),
-            gr.File(label="Other arguments (Optional, YML)", file_types=[".yml", ".yaml"], value=None),
-            gr.Number(
-                label="Samples Per Complex",
-                value=1,
-                minimum=1,
-                maximum=100,
-                precision=0,
-            ),
-            gr.Checkbox(label="Keep Local Structures", value=True),
-            gr.Checkbox(label="Save Visualisation", value=True),
-        ],
-        outputs=gr.File(label="Result"),
-        title="DiffDock-Pocket",
-        description="""
-        Run [DiffDock-Pocket](https://github.com/plainerman/DiffDock-Pocket) for a single protein and ligand.
-        We have provided the most important inputs as UI elements.  
-        Additional values can be included in "Other arguments", and will be passed 
+
+    with gr.Blocks(title="DiffDock-Pocket Web") as demo:
+        gr.Markdown("# DiffDock-Pocket")
+        gr.Markdown("""Run [DiffDock-Pocket](https://github.com/plainerman/DiffDock-Pocket) for a single protein and ligand.
+        We have provided the most important inputs as UI elements.  """)
+        with gr.Row():
+            protein_pdb = gr.File(label="Protein PDB (required)", file_types=[".pdb"])
+            ligand_sdf = gr.File(label="Ligand SDF (required)", file_types=[".sdf"])
+        with gr.Row():
+            samples_per_complex = gr.Number(label="Samples Per Complex", value=1, minimum=1, maximum=100, precision=0)
+            with gr.Column():
+                keep_local_structures = gr.Checkbox(label="Keep Local Structures", value=True)
+                save_vis = gr.Checkbox(label="Save Visualisation", value=True)
+        with gr.Row():
+            gr.Markdown("""Additional values can be included in "Other arguments", and will be passed 
         to [inference.py](https://github.com/plainerman/DiffDock-Pocket/blob/main/inference.py). 
-        Must be a YAML file without any nesting.  
-        For example, to specify a pocket of (0.0, 1.0, 2.0), the YAML file should contain:  
-        ```
-        pocket_center_x: 0.0  
-        pocket_center_y: 1.0  
-        pocket_center_z: 2.0  
-        ```
-        """
-    )
+        Must be a YAML file without any nesting.  """)
+            other_args = gr.File(label="Other arguments (Optional, YML)", file_types=[".yml", ".yaml"], value=None)
+
+        with gr.Row():
+            run_btn = gr.Button("Run DiffDock-Pocket")
+
+        with gr.Row():
+            message = gr.Text(label="Run message", interactive=False)
+        with gr.Row():
+            output_file = gr.File(label="Output Files")
+        with gr.Row():
+            init_value = "Rank1 Reverse Processed Protein will be visualised here"
+            viewer = gr.HTML(value=init_value, label="Protein Viewer", show_label=True)
+
+        _inputs = [protein_pdb, ligand_sdf, other_args, samples_per_complex, keep_local_structures, save_vis]
+        _outputs = [message, output_file, viewer]
+        run_btn.click(fn=run_wrapper, inputs=_inputs, outputs=_outputs)
 
-    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
 
 
 if __name__ == "__main__":

mol_viewer.py

@@ -0,0 +1,64 @@
+viewer_description = """
+PDB viewer using 3Dmol.js  
+If using please cite:  
+3Dmol.js: Molecular visualization with WebGL, Nicholas Rego, David Koes , Bioinformatics, Volume 31, Issue 8, April 2015, Pages 1322–1324, https://doi.org/10.1093/bioinformatics/btu829
+
+See also:
+https://huggingface.co/blog/spaces_3dmoljs
+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>
+            <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>"""
+

run_utils.py

@@ -57,13 +57,20 @@ def kwargs_to_cli_args(**kwargs) -> List[str]:
     return cli_args
 
 
+def read_file_lines(fi_path: str):
+    with open(fi_path, "r") as fp:
+        lines = fp.readlines()
+    mol = "".join(lines)
+    return mol
+
+
 def run_cli_command(
     protein_path: str,
     ligand: str,
     samples_per_complex: int,
     keep_local_structures: bool,
     save_visualisation: bool,
-    other_arg_file: str,
+    other_arg_file: str = None,
     work_dir=None,
 ):
     if work_dir is None:
@@ -98,20 +105,35 @@ def run_cli_command(
 
         # Running the command
         try:
-            result = subprocess.run(
-                command,
-                cwd=work_dir,
-                check=False,
-                text=True,
-                capture_output=True,
-                env=os.environ,
-            )
-            logging.debug(f"Command output:\n{result.stdout}")
-            if result.stderr:
-                logging.error(f"Command error:\n{result.stderr}")
+            skip_running = False
+            if not skip_running:
+                result = subprocess.run(
+                    command,
+                    cwd=work_dir,
+                    check=False,
+                    text=True,
+                    capture_output=True,
+                    env=os.environ,
+                )
+                logging.debug(f"Command output:\n{result.stdout}")
+                if result.stderr:
+                    logging.error(f"Command error:\n{result.stderr}")
         except subprocess.CalledProcessError as e:
             logging.error(f"An error occurred while executing the command: {e}")
 
+        # 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
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -122,7 +144,7 @@ def run_cli_command(
 
         logging.debug(f"Directory '{temp_dir}' zipped to '{full_zip_path}'")
 
-    return full_zip_path
+    return full_zip_path, display_pdb_text
 
 
 if __name__ == "__main__":