File size: 9,208 Bytes
ce9690f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | import collections
import datetime
import logging
import os
from typing import Tuple, Optional, Dict
import gradio as gr
import mol_viewer
import run_utils
from run_utils import PROJECT_URL, PROJECT_DIR, TEMP_DIR
DEFAULT_INFERENCE_ARGS = os.path.join(PROJECT_DIR, "default_inference_args.yaml")
def run_wrapper(protein_pdb_id, protein_file, ligand_smile, ligand_file, config_file, *args) -> Tuple[str, Optional[str], Optional[Dict], Optional[gr.Dropdown]]:
if protein_pdb_id is not None and protein_file is None:
protein_file_name = run_utils.download_pdb(protein_pdb_id, TEMP_DIR)
else:
protein_file_name = protein_file['name']
if protein_file_name is None:
return "Protein file is missing! Must provide a protein file in PDB format", None, None, None
if ligand_file is None and ligand_smile is None:
return "Ligand is missing! Must provide a ligand file in SDF format or SMILE string", None, None, None
config_path = config_file['name'] if config_file else DEFAULT_INFERENCE_ARGS
ligand_desc = ligand_file['name'] if ligand_file else ligand_smile
output_file = run_utils.run_cli_command(
protein_file_name, ligand_desc, config_path, *args,
)
message = f"Calculation completed at {datetime.datetime.now()}"
view_selector_content = collections.OrderedDict()
dropdown = None
# print(f"Output file: {output_file}")
if output_file:
pdb_files, sdf_files = run_utils.process_zip_file(output_file)
# print(f"PDB file: {pdb_files}")
pdb_file = pdb_files[0] if pdb_files else None
for sdf_file in sdf_files:
confidence = sdf_file.get("confidence", None)
# rank1 has no confidence
if confidence is None:
continue
label = f"Rank {sdf_file['rank']}. Confidence {confidence:.2f}"
pdb_text = pdb_file['content'] if pdb_file else None
sdf_text = sdf_file['content']
output_viz = "Output visualisation unavailable"
if pdb_text:
logging.debug(f"Creating 3D visualisation")
output_viz = mol_viewer.gen_3dmol_vis(pdb_text, sdf_text)
view_selector_content[label] = output_viz
labels = list(view_selector_content.keys())
init_value = labels[0] if labels else None
dropdown = gr.Dropdown(interactive=True, label="Ranked samples",
choices=labels, value=init_value)
return message, output_file, view_selector_content, dropdown
def update_view(view_selector_content, view_result_selector, default_str="Output visualisation unavailable"):
if view_selector_content and view_result_selector:
return view_selector_content.get(view_result_selector, default_str)
return default_str
def run():
with gr.Blocks(title="DiffDock Web") as demo:
gr.Markdown("# DiffDock Web")
gr.Markdown(f"""Run [DiffDock]({PROJECT_URL}) for a single protein and ligand.
We have provided the most important inputs as UI elements. """)
with gr.Box():
gr.Markdown("# Input")
with gr.Row():
with gr.Column():
gr.Markdown("## Protein")
protein_pdb_id = gr.Textbox(
placeholder="PDB Code or upload file below", label="Input PDB ID"
)
protein_pdb_file = gr.File(file_count="single", label="Input PDB File")
with gr.Column():
gr.Markdown("## Ligand")
ligand_smile = gr.Textbox(
placeholder="Provide SMILES input or upload mol2/sdf file below",
label="SMILES string",
)
ligand_file = gr.File(file_count="single", label="Input Ligand", file_types=[".sdf", ".mol2"])
with gr.Row():
samples_per_complex = gr.Number(label="Samples Per Complex", value=10, minimum=1, maximum=100, precision=0, interactive=True)
with gr.Row():
with gr.Column():
config_instructions = f"""## Configuration (Optional)
Configuration file to be passed
to [inference.py]({PROJECT_URL}/blob/main/inference.py).
If this is provided, it must supply all necessary arguments.
If not provided, the [default configuration]({PROJECT_URL}/blob/main/app/default_inference_args.yml) will be used."""
gr.Markdown(config_instructions)
config_file = gr.File(label="Configuration (Optional, YML)", file_types=[".yml", ".yaml"], value=None,
info="Additional arguments to pass to DiffDock.")
with gr.Row():
with gr.Column():
gr.Markdown("## Examples")
gr.Examples(
[
[
"6w70",
"examples/6w70.pdb",
"COc1ccc(cc1)n2c3c(c(n2)C(=O)N)CCN(C3=O)c4ccc(cc4)N5CCCCC5=O",
"examples/6w70_ligand.sdf",
10,
True
],
[
"6moa",
"examples/6moa_protein_processed.pdb",
"",
"examples/6moa_ligand.sdf",
10,
True
],
[
"",
"examples/6o5u_protein_processed.pdb",
"",
"examples/6o5u_ligand.sdf",
10,
True
],
[
"",
"examples/6o5u_protein_processed.pdb",
"[NH3+]C[C@H]1O[C@H](O[C@@H]2[C@@H]([NH3+])C[C@H]([C@@H]([C@H]2O)O[C@H]2O[C@H](CO)[C@H]([C@@H]([C@H]2O)[NH3+])O)[NH3+])[C@@H]([C@H]([C@@H]1O)O)O",
"examples/6o5u_ligand.sdf",
10,
True
],
[
"",
"examples/6o5u_protein_processed.pdb",
"",
"examples/6o5u_ligand.sdf",
10,
True
],
[
"",
"examples/6ahs_protein_processed.pdb",
"",
"examples/6ahs_ligand.sdf",
10,
True
],
],
[protein_pdb_id, protein_pdb_file, ligand_smile, ligand_file, samples_per_complex],
)
with gr.Row():
run_btn = gr.Button("Run DiffDock")
with gr.Box():
gr.Markdown("# Output")
with gr.Row():
message = gr.Text(label="Run message", interactive=False)
with gr.Row():
output_file = gr.File(label="Output Files")
with gr.Row():
with gr.Column():
init_value = "DiffDock prediction visualization"
view_result_selector = gr.Dropdown(interactive=True, label="Ranked samples")
viewer = gr.HTML(value=init_value, label="Protein Viewer", show_label=True)
with gr.Row():
gr.Markdown("Many thanks to [Simon Duerr](https://huggingface.co/simonduerr), who created the "
"[original DiffDock web interface](https://huggingface.co/spaces/simonduerr/diffdock), "
"on which this interface is based.")
view_selector_content = gr.Variable()
_inputs = [protein_pdb_id, protein_pdb_file, ligand_smile, ligand_file, config_file]
# See run_utils.py:ARG_ORDER for the order of these arguments
_inputs += [samples_per_complex]
_outputs = [message, output_file, view_selector_content, view_result_selector]
run_btn.click(fn=run_wrapper, inputs=_inputs, outputs=_outputs, preprocess=False)
view_result_selector.change(fn=update_view,
inputs=[view_selector_content, view_result_selector],
outputs=viewer)
server_port = int(os.environ.get("GRADIO_SERVER_PORT", "7860"))
demo.launch(server_name="0.0.0.0", server_port=server_port, share=False)
if __name__ == "__main__":
run_utils.set_env_variables()
run_utils.configure_logging()
run()
|