Spaces:

ronig
/

protein_binding_search

Running

App Files Files Community

roni commited on Mar 13, 2023

Commit

79b6488

1 Parent(s): dfecb5b

visualization enhancement

Browse files

Files changed (2) hide show

app.py +42 -20
protein_viz.py +5 -0

app.py CHANGED Viewed

@@ -8,6 +8,28 @@ 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:
@@ -22,22 +44,10 @@ def format_search_results(raw_search_results):
     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:
@@ -51,21 +61,33 @@ with gr.Blocks() as demo:
             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__":

 engine = get_engine(index_repo, model_repo)
+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)
+    results_options = update_dropdown_menu(search_res)
+    formatted_search_results = format_search_results(search_res)
+    return formatted_search_results, results_options
+def limit_n_results(n):
+    return max(min(n, 20), 1)
+def update_dropdown_menu(search_res):
+    choices = [
+        ','.join([row['pdb_name'], row['chain_id']])
+        for row in search_res
+    ]
+    return gr.Dropdown.update(
+        choices=choices, interactive=True, value=choices[0], visible=True
+    )
 def format_search_results(raw_search_results):
     formatted_search_results = {}
     for res in raw_search_results:
     return formatted_search_results
+def switch_viz(new_choice):
+    pdb_id, chain = new_choice.split(',')
+    title_update = gr.Markdown.update(visible=True)
+    return render_html(pdb_id=pdb_id, chain=chain), title_update
 with gr.Blocks() as demo:
             with gr.Row():
                 with gr.Column():
                     seq_input = gr.Textbox(
+                        value="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")
+            viz_header = gr.Markdown("## Visualization", visible=False)
+            results_selector = gr.Dropdown(
+                choices=[], multiselect=False, visible=False,
+                label="Visualized Search Result"
+            )
             protein_viz = gr.HTML(
+                value=render_html(pdb_id=None, chain=None),
                 label="Protein Visualization"
             )
+            gr.Examples([
+                "KFLIYQMECSTMIFGL",
+                "PHFAMPPIHEDHLE",
+                "AEERIISLD"
+            ], inputs=[seq_input])
     search_button.click(
         search_and_display,
         inputs=[seq_input, n_results],
+        outputs=[search_results, results_selector]
+    )
+    results_selector.change(
+        switch_viz, inputs=results_selector, outputs=[protein_viz, viz_header]
     )
 if __name__ == "__main__":

protein_viz.py CHANGED Viewed

@@ -1,4 +1,6 @@
 def render_html(pdb_id, chain):
     html = f"""
         "<html>
         <header>
@@ -8,7 +10,10 @@ def render_html(pdb_id, chain):
         <body>
         <div style="height: 400px; position: relative;" class="viewer_3Dmoljs"
              data-pdb="{pdb_id}"
              data-select1="chain:{chain}"
              data-style1="cartoon:color=spectrum"
              />
         </body>

 def render_html(pdb_id, chain):
+    if pdb_id is None or chain is None:
+        return ""
     html = f"""
         "<html>
         <header>
         <body>
         <div style="height: 400px; position: relative;" class="viewer_3Dmoljs"
              data-pdb="{pdb_id}"
+             data-backgroundalpha="0.0"
+             data-style="cartoon:color=white"
              data-select1="chain:{chain}"
+             data-zoomto="chain:{chain}"
              data-style1="cartoon:color=spectrum"
              />
         </body>