import gradio as gr import requests from Bio.PDB import PDBParser, MMCIFParser, PDBIO from Bio.PDB.Polypeptide import is_aa from Bio.SeqUtils import seq1 from typing import Optional, Tuple import numpy as np import os from gradio_molecule3d import Molecule3D from model_loader import load_model import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader import re import pandas as pd import copy import transformers, datasets from transformers import AutoTokenizer from transformers import DataCollatorForTokenClassification from datasets import Dataset from scipy.special import expit # Load model and move to device checkpoint = 'ThorbenF/prot_t5_xl_uniref50' max_length = 1500 model, tokenizer = load_model(checkpoint, max_length) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.to(device) model.eval() def normalize_scores(scores): min_score = np.min(scores) max_score = np.max(scores) return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores def read_mol(pdb_path): """Read PDB file and return its content as a string""" with open(pdb_path, 'r') as f: return f.read() def fetch_structure(pdb_id: str, output_dir: str = ".") -> Optional[str]: """ Fetch the structure file for a given PDB ID. Prioritizes CIF files. If a structure file already exists locally, it uses that. """ file_path = download_structure(pdb_id, output_dir) if file_path: return file_path else: return None def download_structure(pdb_id: str, output_dir: str) -> Optional[str]: """ Attempt to download the structure file in CIF or PDB format. Returns the path to the downloaded file, or None if download fails. """ for ext in ['.cif', '.pdb']: file_path = os.path.join(output_dir, f"{pdb_id}{ext}") if os.path.exists(file_path): return file_path url = f"https://files.rcsb.org/download/{pdb_id}{ext}" try: response = requests.get(url, timeout=10) if response.status_code == 200: with open(file_path, 'wb') as f: f.write(response.content) return file_path except Exception as e: print(f"Download error for {pdb_id}{ext}: {e}") return None def convert_cif_to_pdb(cif_path: str, output_dir: str = ".") -> str: """ Convert a CIF file to PDB format using BioPython and return the PDB file path. """ pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb')) parser = MMCIFParser(QUIET=True) structure = parser.get_structure('protein', cif_path) io = PDBIO() io.set_structure(structure) io.save(pdb_path) return pdb_path def fetch_pdb(pdb_id): pdb_path = fetch_structure(pdb_id) if not pdb_path: return None _, ext = os.path.splitext(pdb_path) if ext == '.cif': pdb_path = convert_cif_to_pdb(pdb_path) return pdb_path def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list) -> str: """ Create a PDB file with only the specified chain and replace B-factor with prediction scores """ # Read the original PDB file parser = PDBParser(QUIET=True) structure = parser.get_structure('protein', input_pdb) # Prepare a new structure with only the specified chain new_structure = structure.copy() for model in new_structure: # Remove all chains except the specified one chains_to_remove = [chain for chain in model if chain.id != chain_id] for chain in chains_to_remove: model.detach_child(chain.id) # Create a modified PDB with scores in B-factor scores_dict = {resi: score for resi, score in residue_scores} for model in new_structure: for chain in model: for residue in chain: if residue.id[1] in scores_dict: for atom in residue: atom.bfactor = scores_dict[residue.id[1]] #* 100 # Scale score to B-factor range # Save the modified structure output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_scored.pdb" io = PDBIO() io.set_structure(new_structure) io.save(output_pdb) return output_pdb def calculate_geometric_center(pdb_path: str, high_score_residues: list, chain_id: str): """ Calculate the geometric center of high-scoring residues """ parser = PDBParser(QUIET=True) structure = parser.get_structure('protein', pdb_path) # Collect coordinates of CA atoms from high-scoring residues coords = [] for model in structure: for chain in model: if chain.id == chain_id: for residue in chain: if residue.id[1] in high_score_residues: if 'CA' in residue: # Use alpha carbon as representative ca_atom = residue['CA'] coords.append(ca_atom.coord) # Calculate geometric center if coords: center = np.mean(coords, axis=0) return center return None def process_pdb(pdb_id_or_file, segment): # Determine if input is a PDB ID or file path if pdb_id_or_file.endswith('.pdb'): pdb_path = pdb_id_or_file pdb_id = os.path.splitext(os.path.basename(pdb_path))[0] else: pdb_id = pdb_id_or_file pdb_path = fetch_pdb(pdb_id) if not pdb_path: return "Failed to fetch PDB file", None, None # Determine the file format and choose the appropriate parser _, ext = os.path.splitext(pdb_path) parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True) try: # Parse the structure file structure = parser.get_structure('protein', pdb_path) except Exception as e: return f"Error parsing structure file: {e}", None, None # Extract the specified chain try: chain = structure[0][segment] except KeyError: return "Invalid Chain ID", None, None protein_residues = [res for res in chain if is_aa(res)] sequence = "".join(seq1(res.resname) for res in protein_residues) sequence_id = [res.id[1] for res in protein_residues] # Prepare input for model prediction input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device) with torch.no_grad(): outputs = model(input_ids).logits.detach().cpu().numpy().squeeze() # Calculate scores and normalize them scores = expit(outputs[:, 1] - outputs[:, 0]) normalized_scores = normalize_scores(scores) # Zip residues with scores to track the residue ID and score residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)] # Identify high and mid scoring residues high_score_residues = [resi for resi, score in residue_scores if score > 0.75] mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75] # Calculate geometric center of high-scoring residues geo_center = calculate_geometric_center(pdb_path, high_score_residues, segment) pymol_selection = f"select high_score_residues, resi {'+'.join(map(str, high_score_residues))} and chain {segment}" pymol_center_cmd = f"show spheres, resi {'+'.join(map(str, high_score_residues))} and chain {segment}" if geo_center is not None else "" # Generate the result string current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n" result_str += "Columns: Residue Name, Residue Number, One-letter Code, Normalized Score\n\n" result_str += "\n".join([ f"{res.resname} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" for i, res in enumerate(protein_residues)]) # Create prediction and scored PDB files prediction_file = f"{pdb_id}_predictions.txt" with open(prediction_file, "w") as f: f.write(result_str) # Create chain-specific PDB with scores in B-factor scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores) # Molecule visualization with updated script mol_vis = molecule(pdb_path, residue_scores, segment) # Construct PyMOL command suggestions pymol_commands = f""" PyMOL Visualization Commands: 1. Load PDB: load {os.path.abspath(pdb_path)} 2. Select high-scoring residues: {pymol_selection} 3. Highlight high-scoring residues: show sticks, high_score_residues {pymol_center_cmd} """ return result_str + "\n\n" + pymol_commands, mol_vis, [prediction_file, scored_pdb] def molecule(input_pdb, residue_scores=None, segment='A'): mol = read_mol(input_pdb) # Read PDB file content # Prepare high-scoring residues script if scores are provided high_score_script = "" if residue_scores is not None: # Filter residues based on their scores high_score_residues = [resi for resi, score in residue_scores if score > 0.75] mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75] high_score_script = """ // Load the original model and apply white cartoon style let chainModel = viewer.addModel(pdb, "pdb"); chainModel.setStyle({}, {}); chainModel.setStyle( {"chain": "%s"}, {"cartoon": {"color": "white"}} ); // Create a new model for high-scoring residues and apply red sticks style let highScoreModel = viewer.addModel(pdb, "pdb"); highScoreModel.setStyle({}, {}); highScoreModel.setStyle( {"chain": "%s", "resi": [%s]}, {"stick": {"color": "red"}} ); // Create a new model for medium-scoring residues and apply orange sticks style let midScoreModel = viewer.addModel(pdb, "pdb"); midScoreModel.setStyle({}, {}); midScoreModel.setStyle( {"chain": "%s", "resi": [%s]}, {"stick": {"color": "orange"}} ); """ % ( segment, segment, ", ".join(str(resi) for resi in high_score_residues), segment, ", ".join(str(resi) for resi in mid_score_residues) ) # Generate the full HTML content html_content = f"""
""" # Return the HTML content within an iframe safely encoded for special characters return f'' # Gradio UI with gr.Blocks() as demo: gr.Markdown("# Protein Binding Site Prediction") with gr.Row(): pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...") visualize_btn = gr.Button("Visualize Structure") molecule_output2 = Molecule3D(label="Protein Structure", reps=[ { "model": 0, "style": "cartoon", "color": "whiteCarbon", "residue_range": "", "around": 0, "byres": False, } ]) with gr.Row(): segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...") prediction_btn = gr.Button("Predict Binding Site") molecule_output = gr.HTML(label="Protein Structure") predictions_output = gr.Textbox(label="Binding Site Predictions") download_output = gr.File(label="Download Files", file_count="multiple") prediction_btn.click( process_pdb, inputs=[ pdb_input, segment_input ], outputs=[predictions_output, molecule_output, download_output] ) visualize_btn.click( fetch_pdb, inputs=[pdb_input], outputs=molecule_output2 ) gr.Markdown("## Examples") gr.Examples( examples=[ ["7RPZ", "A"], ["2IWI", "B"], ["2F6V", "A"] ], inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output] ) demo.launch(share=True)