DiffLinker

Runtime error

App Files Files Community

igashov commited on Oct 16, 2022

Commit

52bf9df

1 Parent(s): f9310fd

update logic

Browse files

Files changed (2) hide show

app.py +45 -54
output.py +98 -0

app.py CHANGED Viewed

@@ -3,6 +3,7 @@ import numpy as np
 import os
 import torch
 import subprocess
 from rdkit import Chem
 from src import const
@@ -12,46 +13,6 @@ from src.lightning import DDPM
 from src.linker_size_lightning import SizeClassifier
-HTML_TEMPLATE = """<!DOCTYPE html>
-<html>
-    <head>
-        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
-        <style>
-            .mol-container {{
-                width: 600px;
-                height: 600px;
-                position: relative;
-            }}
-            .mol-container select{{
-                background-image:None;
-            }}
-        </style>
-        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
-    </head>
-    <body>
-    <div id="container" class="mol-container"></div>
-    <script>
-        $(document).ready(function() {{
-            let element = $("#container");
-            let config = {{ backgroundColor: "white" }};
-            let viewer = $3Dmol.createViewer( element, config );
-            viewer.addModel(`{molecule}`, "{fmt}")
-            viewer.getModel().setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
-            viewer.zoomTo();
-            viewer.render();
-        }});
-    </script>
-    </body>
-</html>
-"""
-IFRAME_TEMPLATE = """<iframe style="width: 100%; height: 700px" 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='{html}'></iframe>"""
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 os.makedirs("results", exist_ok=True)
 os.makedirs("models", exist_ok=True)
@@ -100,17 +61,42 @@ def read_molecule(path):
     raise Exception('Unknown file extension')
 def generate(input_file):
     try:
-        path = input_file.name
         molecule = read_molecule(path)
         name = '.'.join(path.split('/')[-1].split('.')[:-1])
         inp_sdf = f'results/{name}_input.sdf'
         inp_xyz = f'results/{name}_input.xyz'
         out_sdf = f'results/{name}_output.sdf'
         out_xyz = f'results/{name}_output.xyz'
-        print(f'Input path={path}, name={name}')
     except Exception as e:
         return f'Could not read the molecule: {e}'
@@ -153,9 +139,9 @@ def generate(input_file):
         break
     generated_molecule = read_molecule_content(out_sdf)
-    html = HTML_TEMPLATE.format(molecule=generated_molecule, fmt='sdf')
     return [
-        IFRAME_TEMPLATE.format(html=html),
         [inp_sdf, inp_xyz, out_sdf, out_xyz],
     ]
@@ -167,16 +153,21 @@ with demo:
         with gr.Row():
             with gr.Column():
                 gr.Markdown('## Input Fragments')
-                gr.Markdown('Upload the file with 3D-coordinates of the input fragments in .pdb, .mol2 or .sdf format')
                 input_file = gr.File(file_count='single', label='Input fragments')
-    button = gr.Button('Generate Linker!')
-    gr.Markdown('')
-    gr.Markdown('## Output')
-    visualization = gr.HTML()
-    output_files = gr.File(file_count="multiple", label="Output Files")
     button.click(
         fn=generate,
         inputs=[input_file],

 import os
 import torch
 import subprocess
+import output
 from rdkit import Chem
 from src import const
 from src.linker_size_lightning import SizeClassifier
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 os.makedirs("results", exist_ok=True)
 os.makedirs("models", exist_ok=True)
     raise Exception('Unknown file extension')
+def show_input(input_file):
+    if input_file is None:
+        return ''
+    path = input_file.name
+    extension = path.split('.')[-1]
+    if extension not in ['sdf', 'pdb', 'mol', 'mol2']:
+        msg = output.INVALID_FORMAT_MSG.format(extension=extension)
+        return output.IFRAME_TEMPLATE.format(html=msg)
+    try:
+        molecule = read_molecule_content(path)
+    except Exception as e:
+        return f'Could not read the molecule: {e}'
+    html = output.HTML_TEMPLATE.format(molecule=molecule, fmt=extension)
+    return output.IFRAME_TEMPLATE.format(html=html)
 def generate(input_file):
+    if input_file is None:
+        return ''
+    path = input_file.name
+    extension = path.split('.')[-1]
+    if extension not in ['sdf', 'pdb', 'mol', 'mol2']:
+        msg = output.INVALID_FORMAT_MSG.format(extension=extension)
+        return output.IFRAME_TEMPLATE.format(html=msg)
     try:
         molecule = read_molecule(path)
         name = '.'.join(path.split('/')[-1].split('.')[:-1])
         inp_sdf = f'results/{name}_input.sdf'
         inp_xyz = f'results/{name}_input.xyz'
         out_sdf = f'results/{name}_output.sdf'
         out_xyz = f'results/{name}_output.xyz'
     except Exception as e:
         return f'Could not read the molecule: {e}'
         break
     generated_molecule = read_molecule_content(out_sdf)
+    html = output.HTML_TEMPLATE.format(molecule=generated_molecule, fmt='sdf')
     return [
+        output.IFRAME_TEMPLATE.format(html=html),
         [inp_sdf, inp_xyz, out_sdf, out_xyz],
     ]
         with gr.Row():
             with gr.Column():
                 gr.Markdown('## Input Fragments')
+                gr.Markdown('Upload the file with 3D-coordinates of the input fragments in .pdb, .mol2 or .sdf format:')
                 input_file = gr.File(file_count='single', label='Input fragments')
+                button = gr.Button('Generate Linker!')
+                gr.Markdown('')
+                gr.Markdown('## Output Files')
+                gr.Markdown('Download files with the generated molecules here:')
+                output_files = gr.File(file_count='multiple', label='Output Files')
+            with gr.Column():
+                visualization = gr.HTML()
+    input_file.change(
+        fn=show_input,
+        inputs=[input_file],
+        outputs=[visualization],
+    )
     button.click(
         fn=generate,
         inputs=[input_file],

output.py ADDED Viewed

	@@ -0,0 +1,98 @@

+HTML_TEMPLATE = """<!DOCTYPE html>
+<html>
+    <head>
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <style>
+            .mol-container {{
+                width: 600px;
+                height: 600px;
+                position: relative;
+            }}
+            .mol-container select{{
+                background-image:None;
+            }}
+        </style>
+        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>
+    <div id="container" class="mol-container"></div>
+    <script>
+        $(document).ready(function() {{
+            let element = $("#container");
+            let config = {{ backgroundColor: "white" }};
+            let viewer = $3Dmol.createViewer( element, config );
+            viewer.addModel(`{molecule}`, "{fmt}")
+            viewer.getModel().setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
+            viewer.zoomTo();
+            viewer.render();
+        }});
+    </script>
+    </body>
+</html>
+"""
+INVALID_FORMAT_MSG = """
+<!DOCTYPE html>
+<html>
+    <head>
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <style>
+            body{{
+                font-family:sans-serif
+            }}
+        </style>
+    </head>
+    <body>
+        <h3>Invalid file format: {extension}</h3>
+        Please upload the file in one of the following formats:
+        <ul>
+          <li>.pdb</li>
+          <li>.sdf</li>
+          <li>.mol</li>
+          <li>.mol2</li>
+        </ul>
+    </body>
+</html>
+"""
+HTML_TEMPLATE = """<!DOCTYPE html>
+<html>
+    <head>
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <style>
+            .mol-container {{
+                width: 600px;
+                height: 600px;
+                position: relative;
+            }}
+            .mol-container select{{
+                background-image:None;
+            }}
+        </style>
+        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>
+    <div id="container" class="mol-container"></div>
+    <script>
+        $(document).ready(function() {{
+            let element = $("#container");
+            let config = {{ backgroundColor: "white" }};
+            let viewer = $3Dmol.createViewer( element, config );
+            viewer.addModel(`{molecule}`, "{fmt}")
+            viewer.getModel().setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
+            viewer.zoomTo();
+            viewer.render();
+        }});
+    </script>
+    </body>
+</html>
+"""
+IFRAME_TEMPLATE = """<iframe style="width: 100%; height: 700px" 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='{html}'></iframe>"""