app.py

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"""
    <!DOCTYPE html>
    <html>
    <head>    
        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
        <style>
        .mol-container {{
            width: 100%;
            height: 700px;
            position: relative;
        }}
        </style>
        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js"></script>
        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
    </head>
    <body>
        <div id="container" class="mol-container"></div>
        <script>
            let pdb = `{mol}`; // Use template literal to properly escape PDB content
            $(document).ready(function () {{
                let element = $("#container");
                let config = {{ backgroundColor: "white" }};
                let viewer = $3Dmol.createViewer(element, config);
                
                {high_score_script}
                
                // Add hover functionality
                viewer.setHoverable(
                    {{}}, 
                    true, 
                    function(atom, viewer, event, container) {{
                        if (!atom.label) {{
                            atom.label = viewer.addLabel(
                                atom.resn + ":" +atom.resi + ":" + atom.atom, 
                                {{
                                    position: atom, 
                                    backgroundColor: 'mintcream', 
                                    fontColor: 'black',
                                    fontSize: 12,
                                    padding: 2
                                }}
                            );
                        }}
                    }},
                    function(atom, viewer) {{
                        if (atom.label) {{
                            viewer.removeLabel(atom.label);
                            delete atom.label;
                        }}
                    }}
                );
                
                viewer.zoomTo();
                viewer.render();
                viewer.zoom(0.8, 2000);
            }});
        </script>
    </body>
    </html>
    """
    
    # Return the HTML content within an iframe safely encoded for special characters
    return f'<iframe width="100%" height="700" srcdoc="{html_content.replace(chr(34), "&quot;").replace(chr(39), "&#39;")}"></iframe>'


# 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)