Spaces:

ronig
/

protein_binding_search

Running

App Files Files Community

roni commited on Mar 13, 2023

Commit

dfecb5b

1 Parent(s): 20d73a0

visualizigin search results

Browse files

Files changed (3) hide show

app.py +44 -18
get_index.py +1 -1
protein_viz.py +21 -0

app.py CHANGED Viewed

@@ -1,32 +1,45 @@
 import gradio as gr
 from get_index import get_engine
 index_repo = "ronig/siamese_protein_index"
 model_repo = "ronig/protein_search_engine"
 engine = get_engine(index_repo, model_repo)
-def search(seq, n_res):
-    n_res = int(limit_n_results(n_res))
-    search_results = engine.search_by_sequence(seq, n=n_res)
-    outputs = {}
-    for res in search_results:
         prot = res["pdb_name"]
         chain = res["chain_id"]
         value = res["score"]
-        genes = ','.join(res["gene_ids"])
         key = f"PDB: {prot} | Chain: {chain}"
-        if genes != 'Unknown':
             key += f" | Gene: {genes}"
-        outputs[key] = value
-    return outputs
 def limit_n_results(n):
     return max(min(n, 20), 1)
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("""
@@ -34,13 +47,26 @@ with gr.Blocks() as demo:
         This application examines all files uploaded to [PDB](https://www.rcsb.org/) to find the chains with which a given protein sequence is most likely to bind.
         If the results are linked to a specific gene, their IDs will also be displayed.
         """)
-        with gr.Row():
-            with gr.Column():
-                seq_input = gr.Textbox("KFLIYQMECSTMIFGL", label="Input Sequence")
-                n_results = gr.Number(5, label="N Results")
-                search_button = gr.Button("Search")
-            output = gr.Label(num_top_classes=20, label="Search Results")
-    search_button.click(search, inputs=[seq_input, n_results], outputs=output)
-if __name__ == '__main__':
     demo.launch()

 import gradio as gr
 from get_index import get_engine
+from protein_viz import render_html
 index_repo = "ronig/siamese_protein_index"
 model_repo = "ronig/protein_search_engine"
 engine = get_engine(index_repo, model_repo)
+def format_search_results(raw_search_results):
+    formatted_search_results = {}
+    for res in raw_search_results:
         prot = res["pdb_name"]
         chain = res["chain_id"]
         value = res["score"]
+        genes = ",".join(res["gene_ids"])
         key = f"PDB: {prot} | Chain: {chain}"
+        if genes != "Unknown":
             key += f" | Gene: {genes}"
+        formatted_search_results[key] = value
+    return formatted_search_results
 def limit_n_results(n):
     return max(min(n, 20), 1)
+def update_html(search_res):
+    first_entry = search_res[0]
+    pdb_id = first_entry['pdb_name']
+    return render_html(pdb_id=pdb_id, chain=first_entry['chain_id'])
+def search_and_display(seq, n_res):
+    n_res = int(limit_n_results(n_res))
+    search_res = engine.search_by_sequence(seq, n=n_res)
+    formatted_search_results = format_search_results(search_res)
+    new_html = update_html(search_res)
+    return formatted_search_results, new_html
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("""
         This application examines all files uploaded to [PDB](https://www.rcsb.org/) to find the chains with which a given protein sequence is most likely to bind.
         If the results are linked to a specific gene, their IDs will also be displayed.
         """)
+        with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    seq_input = gr.Textbox(
+                        placeholder="KFLIYQMECSTMIFGL",
+                        label="Input Sequence"
+                    )
+                    n_results = gr.Number(5, label="N Results")
+                    search_button = gr.Button("Search", variant='primary')
+                search_results = gr.Label(num_top_classes=20, label="Search Results")
+            protein_viz = gr.HTML(
+                value=render_html(pdb_id="2POR", chain='A'),
+                label="Protein Visualization"
+            )
+            gr.Examples(["KFLIYQMECSTMIFGL"], inputs=[seq_input])
+    search_button.click(
+        search_and_display,
+        inputs=[seq_input, n_results],
+        outputs=[search_results, protein_viz]
+    )
+if __name__ == "__main__":
     demo.launch()

get_index.py CHANGED Viewed

@@ -19,5 +19,5 @@ def get_engine(index_repo: str, model_repo: str):
     )
     sys.path.append(str(local_arch_path))
     from protein_index import ProteinSearchEngine  # pylint: disable=import-error,import-outside-toplevel
     return ProteinSearchEngine(data_path=index_path)

     )
     sys.path.append(str(local_arch_path))
     from protein_index import ProteinSearchEngine  # pylint: disable=import-error,import-outside-toplevel
+    index_path = Path("/data/indexes/pdb")
     return ProteinSearchEngine(data_path=index_path)

protein_viz.py ADDED Viewed

	@@ -0,0 +1,21 @@

+def render_html(pdb_id, chain):
+    html = f"""
+        "<html>
+        <header>
+            <script src="https://3Dmol.org/build/3Dmol-min.js"></script>
+            <script src="https://3Dmol.org/build/3Dmol.ui-min.js"></script>
+        </header>
+        <body>
+        <div style="height: 400px; position: relative;" class="viewer_3Dmoljs"
+             data-pdb="{pdb_id}"
+             data-select1="chain:{chain}"
+             data-style1="cartoon:color=spectrum"
+             />
+        </body>
+        </html>
+    """
+    iframe = f"""
+            <iframe style="width: 100%; height: 480px" name="protein-vis"
+            frameborder="0" srcdoc='{html}'></iframe>
+        """
+    return iframe