showing pdb title on visualization

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 
 from get_index import get_engine
-from protein_viz import render_html
+from protein_viz import get_protein_name, render_html
 
 index_repo = "ronig/siamese_protein_index"
 model_repo = "ronig/protein_search_engine"
@@ -47,7 +47,12 @@ def format_search_results(raw_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
+    protein_name = get_protein_name(pdb_id)
+    description_update = gr.Markdown.update(
+        value=f"""**PDB Title**: {protein_name}""",
+        visible=True
+    )
+    return render_html(pdb_id=pdb_id, chain=chain), title_update, description_update
 
 
 with gr.Blocks() as demo:
@@ -72,6 +77,7 @@ with gr.Blocks() as demo:
                 choices=[], multiselect=False, visible=False,
                 label="Visualized Search Result"
             )
+            viz_body = gr.Markdown("", visible=False)
             protein_viz = gr.HTML(
                 value=render_html(pdb_id=None, chain=None),
                 label="Protein Visualization"
@@ -87,7 +93,9 @@ with gr.Blocks() as demo:
         outputs=[search_results, results_selector]
     )
     results_selector.change(
-        switch_viz, inputs=results_selector, outputs=[protein_viz, viz_header]
+        switch_viz,
+        inputs=results_selector,
+        outputs=[protein_viz, viz_header, viz_body]
     )
 
 demo.launch()

protein_viz.py

@@ -1,3 +1,9 @@
+from typing import List
+
+import mygene
+import requests
+
+
 def render_html(pdb_id, chain):
     if pdb_id is None or chain is None:
         return ""
@@ -25,3 +31,21 @@ def render_html(pdb_id, chain):
             frameborder="0" srcdoc='{html}'></iframe>
         """
     return iframe
+
+
+def get_gene_names(genes: List[str]):
+    mg = mygene.MyGeneInfo()
+    ginfo = mg.querymany(genes, scopes='ensembl.gene')
+    gene_names = [gene['name'] for gene in ginfo]
+    return gene_names
+
+
+def get_protein_name(pdb_id: str):
+    url = f"https://data.rcsb.org/rest/v1/core/entry/{pdb_id}"
+    response = requests.get(url, timeout=1)
+    if response.ok:
+        data = response.json()
+        protein_name = data['struct']['title']
+    else:
+        protein_name = 'Unknown'
+    return protein_name

requirements.txt

@@ -1,3 +1,4 @@
 torch
 transformers
 annoy
+mygene