Update app.py

app.py

@@ -5,7 +5,12 @@ from shutil import copyfile
 import gradio as gr
 import os
 
-def get_pdb(pdb_code="", filepath=""):
+def get_pdb(upload_choice="PDB Code", pdb_code="", filepath=""):
+    urls = {
+        "PDB Code": {base:"https://files.rcsb.org/view/", suffix:".pdb"}, 
+        "AlphaFold DB": {base:"https://alphafold.ebi.ac.uk/files/", suffix: "-F1-model_v4.pdb"} 
+        "ESM Atlas": {base: "https://api.esmatlas.com/fetchPredictedStructure/", suffix:".pdb"}
+    }
     if pdb_code is None or pdb_code == "":
         try:
             #move file to home folder to have it accessible from the web
@@ -14,8 +19,8 @@ def get_pdb(pdb_code="", filepath=""):
         except AttributeError as e:
             return None
     else:
-        os.system(f"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb")
-        return f"{pdb_code}.pdb"
+        os.system(f"wget -qnc {urls[upload_choice]['base']}{pdb_code}{urls[upload_choice]['suffix']}")
+        return f"{pdb_code}{urls[upload_choice]['suffix']}"
 
 
 def read_mol(molpath):
@@ -103,10 +108,17 @@ def molecule(input_pdb, public_link):
     allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
 
 
-def update(inp, file, public_link):
-    pdb_path = get_pdb(inp, file)
+def update(upload_choice, inp, file, public_link):
+    pdb_path = get_pdb(upload_choice, inp, file)
     return molecule(pdb_path, public_link)
 
+def toggle_upload_input(choice):
+    if choice not "local file":
+        return gr.update(visible=True, value=None, placeholder=choice), gr.update(visible=False, value=None)
+    elif choice == "local file":
+        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
+
+        
 
 demo = gr.Blocks()
 
@@ -117,12 +129,18 @@ with demo:
     public_link = gr.Variable(value="https://simonduerr-molstar-gradio.hf.space")
     with gr.Row():
         with gr.Box():
+            upload_choice = gr.Radio(["PDB Code", "AlphaFold DB", "ESM Atlas","local file"], label="File source", value='PDB Code')
             inp = gr.Textbox(
-                placeholder="PDB Code or upload file below", label="Input structure"
+                placeholder="PDB Code", label="Input structure"
             )
             file = gr.File(file_count="single")
-            gr.Examples(["2CBA", "6VXX"], inp)
+            upload_choice.change(fn=toggle_upload_input,
+                                    inputs=[upload_choice],
+                                    outputs=[inp, file],
+                                    queue=False)
+
+            gr.Examples([["PDB code", "2CBA"], ["ESM Atlas", "MGYP001531319262"]], [upload_choice,inp])
             btn = gr.Button("View structure")
     mol = gr.HTML()
-    btn.click(fn=update, inputs=[inp, file, public_link], outputs=mol)
+    btn.click(fn=update, inputs=[upload_choice, inp, file, public_link], outputs=mol)
 _, _, pl = demo.launch()  # use public link with share=True locally