Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| from .blocks import upload_pdb_button | |
| from utils.downloader import download_pdb, download_af2 | |
| root_dir = __file__.rsplit("/", 3)[0] | |
| structure_types = ["AlphaFoldDB", "PDB"] | |
| def upload_structure(file: str): | |
| return file | |
| def get_structure_path(structure: str, structure_type: str) -> str: | |
| # If the structure is manually uploaded | |
| if structure[0] == "/": | |
| return structure | |
| # If the structure is a Uniprot ID, download the structure from AlphaFoldDB | |
| elif structure_type == "AlphaFoldDB": | |
| save_path = f"{root_dir}/demo/cache/{structure}.pdb" | |
| if not os.path.exists(save_path): | |
| download_af2(structure, "pdb", save_path) | |
| return save_path | |
| # If the structure is a PDB ID, download the structure from PDB | |
| elif structure_type == "PDB": | |
| save_path = f"{root_dir}/demo/cache/{structure}.cif" | |
| if not os.path.exists(save_path): | |
| download_pdb(structure, "cif", save_path) | |
| return save_path | |
| def tmalign(structure_1: str, structure_type_1: str, structure_2: str, structure_type_2: str): | |
| structure_path_1 = get_structure_path(structure_1, structure_type_1) | |
| structure_path_2 = get_structure_path(structure_2, structure_type_2) | |
| cmd = f"/tmp/TMalign {structure_path_1} {structure_path_2}" | |
| r = os.popen(cmd) | |
| text = r.read() | |
| return text | |
| # Build the block for computing protein-text similarity | |
| def build_TMalign(): | |
| gr.Markdown(f"# Calculate TM-score between two protein structures") | |
| with gr.Row(equal_height=True): | |
| with gr.Column(): | |
| # Compute similarity score between sequence and text | |
| with gr.Row(): | |
| structure_1 = gr.Textbox(label="Protein structure 1 (input Uniprot ID or PDB ID or upload a pdb file)") | |
| structure_type_1 = gr.Dropdown(structure_types, label="Structure type (if the structure is manually uploaded, ignore this field)", | |
| value="AlphaFoldDB", interactive=True, visible=True) | |
| # Provide an upload button to upload a pdb file | |
| upload_btn_1, _ = upload_pdb_button(visible=True, chain_visible=False) | |
| upload_btn_1.upload(upload_structure, inputs=[upload_btn_1], outputs=[structure_1]) | |
| with gr.Row(): | |
| structure_2 = gr.Textbox(label="Protein structure 2 (input Uniprot ID or PDB ID or upload a pdb file)") | |
| structure_type_2 = gr.Dropdown(structure_types, label="Structure type (if the structure is manually uploaded, ignore this field)", | |
| value="AlphaFoldDB", interactive=True, visible=True) | |
| # Provide an upload button to upload a pdb file | |
| upload_btn_2, _ = upload_pdb_button(visible=True, chain_visible=False) | |
| upload_btn_2.upload(upload_structure, inputs=[upload_btn_2], outputs=[structure_2]) | |
| compute_btn = gr.Button(value="Compute TM-score") | |
| tmscore = gr.TextArea(label="TM-score", interactive=False) | |
| compute_btn.click(tmalign, inputs=[structure_1, structure_type_1, structure_2, structure_type_2], | |
| outputs=[tmscore]) | |