app.py

from datetime import datetime
import gradio as gr
import requests
from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select
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
from transformers import AutoTokenizer, DataCollatorForTokenClassification

from datasets import Dataset

from scipy.special import expit

# Load model and move to device
#checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
#checkpoint = 'ThorbenF/prot_t5_xl_uniref50_cryptic'
#checkpoint = 'ThorbenF/prot_t5_xl_uniref50_database'
checkpoint = 'ThorbenF/prot_t5_xl_uniref50_full'
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 = ".") -> 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)
    return file_path

def download_structure(pdb_id: str, output_dir: str) -> str:
    """
    Attempt to download the structure file in CIF or PDB format.
    Returns the path to the downloaded file.
    """
    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}"
        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
    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)
    _, 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, protein_residues: list) -> str:
    """
    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores
    """
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure('protein', input_pdb)
    
    output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb"
    
    # Create scores dictionary for easy lookup
    scores_dict = {resi: score for resi, score in residue_scores}

    # Create a custom Select class
    class ResidueSelector(Select):
        def __init__(self, chain_id, selected_residues, scores_dict):
            self.chain_id = chain_id
            self.selected_residues = selected_residues
            self.scores_dict = scores_dict
        
        def accept_chain(self, chain):
            return chain.id == self.chain_id
        
        def accept_residue(self, residue):
            return residue.id[1] in self.selected_residues

        def accept_atom(self, atom):
            if atom.parent.id[1] in self.scores_dict:
                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100
            return True

    # Prepare output PDB with selected chain and residues, modified B-factors
    io = PDBIO()
    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)
    
    io.set_structure(structure[0])
    io.save(output_pdb, selector)
    
    return output_pdb

def generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, score_type):
    """Generate PyMOL commands based on score type"""
    pymol_commands = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\nScore Type: {score_type}\n\n"
    
    pymol_commands += f"""
    # PyMOL Visualization Commands
    fetch {pdb_id}, protein
    hide everything, all
    show cartoon, chain {segment}
    color white, chain {segment}
    """
    
    # Define colors for each score bracket
    bracket_colors = {
        "0.0-0.2": "white",
        "0.2-0.4": "lightorange",
        "0.4-0.6": "yelloworange",
        "0.6-0.8": "orange",
        "0.8-1.0": "red"
    }
    
    # Add PyMOL commands for each score bracket
    for bracket, residues in residues_by_bracket.items():
        if residues:  # Only add commands if there are residues in this bracket
            color = bracket_colors[bracket]
            resi_list = '+'.join(map(str, residues))
            pymol_commands += f"""
    select bracket_{bracket.replace('.', '').replace('-', '_')}, resi {resi_list} and chain {segment}
    show sticks, bracket_{bracket.replace('.', '').replace('-', '_')}
    color {color}, bracket_{bracket.replace('.', '').replace('-', '_')}
    """
    return pymol_commands

def generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, scores, current_time, score_type):
    """Generate results text based on score type"""
    result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\nScore Type: {score_type}\n\n"
    result_str += "Residues by Score Brackets:\n\n"
    
    # Add residues for each bracket
    for bracket, residues in residues_by_bracket.items():
        result_str += f"Bracket {bracket}:\n"
        result_str += f"Columns: Residue Name, Residue Number, One-letter Code, {score_type} Score\n"
        result_str += "\n".join([
            f"{res.resname} {res.id[1]} {sequence[i]} {scores[i]:.2f}" 
            for i, res in enumerate(protein_residues) if res.id[1] in residues
        ])
        result_str += "\n\n"
    
    return result_str




def process_pdb(pdb_id_or_file, segment, score_type='normalized'):
    # 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)
    
    # 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)
    
    # Parse the structure file
    structure = parser.get_structure('protein', pdb_path)
    
    # Extract the specified chain
    chain = structure[0][segment]
    
    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]

    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
    raw_scores = expit(outputs[:, 1] - outputs[:, 0])
    normalized_scores = normalize_scores(raw_scores)
    
    # Choose which scores to use based on score_type
    display_scores = normalized_scores if score_type == 'normalized' else raw_scores
    
    # Zip residues with scores to track the residue ID and score
    residue_scores = [(resi, score) for resi, score in zip(sequence_id, display_scores)]
    
    # Also save both score types for later use
    raw_residue_scores = [(resi, score) for resi, score in zip(sequence_id, raw_scores)]
    norm_residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]
    
    # Define the score brackets
    score_brackets = {
        "0.0-0.2": (0.0, 0.2),
        "0.2-0.4": (0.2, 0.4),
        "0.4-0.6": (0.4, 0.6),
        "0.6-0.8": (0.6, 0.8),
        "0.8-1.0": (0.8, 1.0)
    }
    
    # Initialize a dictionary to store residues by bracket
    residues_by_bracket = {bracket: [] for bracket in score_brackets}
    
    # Categorize residues into brackets
    for resi, score in residue_scores:
        for bracket, (lower, upper) in score_brackets.items():
            if lower <= score < upper:
                residues_by_bracket[bracket].append(resi)
                break
    
    # Generate timestamp
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Generate result text and PyMOL commands based on score type
    display_score_type = "Normalized" if score_type == 'normalized' else "Raw"
    result_str = generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, 
                                      display_scores, current_time, display_score_type)
    pymol_commands = generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, display_score_type)
    
    # Create chain-specific PDB with scores in B-factor
    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)

    # Molecule visualization with updated script with color mapping
    mol_vis = molecule(pdb_path, residue_scores, segment)
    
    # Create prediction file
    prediction_file = f"{pdb_id}_{display_score_type.lower()}_binding_site_residues.txt"
    with open(prediction_file, "w") as f:
        f.write(result_str)
    
    scored_pdb_name = f"{pdb_id}_{segment}_{display_score_type.lower()}_predictions_scores.pdb"
    os.rename(scored_pdb, scored_pdb_name)
    
    return pymol_commands, mol_vis, [prediction_file, scored_pdb_name], raw_residue_scores, norm_residue_scores, pdb_id, segment

def molecule(input_pdb, residue_scores=None, segment='A'):
    # Read PDB file content
    mol = read_mol(input_pdb)
    
    # Prepare high-scoring residues script if scores are provided
    high_score_script = ""
    if residue_scores is not None:
        # Filter residues based on their scores
        class1_score_residues = [resi for resi, score in residue_scores if 0.0 < score <= 0.2]
        class2_score_residues = [resi for resi, score in residue_scores if 0.2 < score <= 0.4]
        class3_score_residues = [resi for resi, score in residue_scores if 0.4 < score <= 0.6]
        class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]
        class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]
        
        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 class1Model = viewer.addModel(pdb, "pdb");
        class1Model.setStyle({}, {});
        class1Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFFFFFF", "opacity": 0.5}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class2Model = viewer.addModel(pdb, "pdb");
        class2Model.setStyle({}, {});
        class2Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFFD580", "opacity": 0.7}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class3Model = viewer.addModel(pdb, "pdb");
        class3Model.setStyle({}, {});
        class3Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFFA500", "opacity": 1}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class4Model = viewer.addModel(pdb, "pdb");
        class4Model.setStyle({}, {});
        class4Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFF4500", "opacity": 1}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class5Model = viewer.addModel(pdb, "pdb");
        class5Model.setStyle({}, {});
        class5Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFF0000", "alpha": 1}}
        );

        """ % (
            segment,
            segment,
            ", ".join(str(resi) for resi in class1_score_residues),
            segment,
            ", ".join(str(resi) for resi in class2_score_residues),
            segment,
            ", ".join(str(resi) for resi in class3_score_residues),
            segment,
            ", ".join(str(resi) for resi in class4_score_residues),
            segment,
            ", ".join(str(resi) for resi in class5_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: 18,
                                    padding: 4
                                }}
                            );
                        }}
                    }},
                    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>'

with gr.Blocks(css="""
    /* Customize Gradio button colors */
    #visualize-btn, #predict-btn {
        background-color: #FF7300; /* Deep orange */
        color: white;
        border-radius: 5px;
        padding: 10px;
        font-weight: bold;
    }
    #visualize-btn:hover, #predict-btn:hover {
        background-color: #CC5C00; /* Darkened orange on hover */
    }
""") as demo:
    gr.Markdown("# Protein Binding Site Prediction")
    
    # Mode selection
    mode = gr.Radio(
        choices=["PDB ID", "Upload File"],
        value="PDB ID",
        label="Input Mode",
        info="Choose whether to input a PDB ID or upload a PDB/CIF file."
    )

    # Input components based on mode
    pdb_input = gr.Textbox(value="2F6V", label="PDB ID", placeholder="Enter PDB ID here...")
    pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
    visualize_btn = gr.Button("Visualize Structure", elem_id="visualize-btn")

    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 (protein)", placeholder="Enter Chain ID here...",
        info="Choose in which chain to predict binding sites.")
        prediction_btn = gr.Button("Predict Binding Site", elem_id="predict-btn")

    # Add score type selector
    score_type = gr.Radio(
        choices=["Normalized Scores", "Raw Scores"],
        value="Normalized Scores",
        label="Score Visualization Type",
        info="Choose which score type to visualize"
    )

    molecule_output = gr.HTML(label="Protein Structure")
    explanation_vis = gr.Markdown("""
    Score dependent colorcoding:
    - 0.0-0.2: white  
    - 0.2–0.4: light orange  
    - 0.4–0.6: yellow orange
    - 0.6–0.8: orange
    - 0.8–1.0: red
    """)
    predictions_output = gr.Textbox(label="Visualize Prediction with PyMol")
    gr.Markdown("### Download:\n- List of predicted binding site residues\n- PDB with score in beta factor column")
    download_output = gr.File(label="Download Files", file_count="multiple")
    
    # Store these as state variables so we can switch between them
    raw_scores_state = gr.State(None)
    norm_scores_state = gr.State(None)
    last_pdb_path = gr.State(None)
    last_segment = gr.State(None)
    last_pdb_id = gr.State(None)
    
    def process_interface(mode, pdb_id, pdb_file, chain_id, score_type_val):
        selected_score_type = 'normalized' if score_type_val == "Normalized Scores" else 'raw'
        
        # First get the actual PDB file path
        if mode == "PDB ID":
            pdb_path = fetch_pdb(pdb_id)  # Get the actual file path
            
            pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_id_result, segment = process_pdb(pdb_path, chain_id, selected_score_type)
            # Store the actual file path, not just the PDB ID
            return pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_path, chain_id, pdb_id_result
        elif mode == "Upload File":
            _, ext = os.path.splitext(pdb_file.name)
            file_path = os.path.join('./', f"{_}{ext}")
            if ext == '.cif':
                pdb_path = convert_cif_to_pdb(file_path)
            else:
                pdb_path = file_path
            
            pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_id_result, segment = process_pdb(pdb_path, chain_id, selected_score_type)
            return pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_path, chain_id, pdb_id_result
    
    def update_visualization_and_files(score_type_val, raw_scores, norm_scores, pdb_path, segment, pdb_id):
        if raw_scores is None or norm_scores is None or pdb_path is None or segment is None or pdb_id is None:
            return None, None, None
        
        # Choose scores based on radio button selection
        selected_score_type = 'normalized' if score_type_val == "Normalized Scores" else 'raw'
        selected_scores = norm_scores if selected_score_type == 'normalized' else raw_scores
        
        # Generate visualization with selected scores
        mol_vis = molecule(pdb_path, selected_scores, segment)
        
        # Generate PyMOL commands and downloadable files
        # Get structure for residue info
        _, ext = os.path.splitext(pdb_path)
        parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)
        structure = parser.get_structure('protein', pdb_path)
        chain = structure[0][segment]
        protein_residues = [res for res in chain if is_aa(res)]
        sequence = "".join(seq1(res.resname) for res in protein_residues)
        
        # Define score brackets
        score_brackets = {
            "0.0-0.2": (0.0, 0.2),
            "0.2-0.4": (0.2, 0.4),
            "0.4-0.6": (0.4, 0.6),
            "0.6-0.8": (0.6, 0.8),
            "0.8-1.0": (0.8, 1.0)
        }
        
        # Initialize a dictionary to store residues by bracket
        residues_by_bracket = {bracket: [] for bracket in score_brackets}
        
        # Categorize residues into brackets
        for resi, score in selected_scores:
            for bracket, (lower, upper) in score_brackets.items():
                if lower <= score < upper:
                    residues_by_bracket[bracket].append(resi)
                    break
        
        # Generate timestamp
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        # Generate result text and PyMOL commands based on score type
        display_score_type = "Normalized" if selected_score_type == 'normalized' else "Raw"
        scores_array = [score for _, score in selected_scores]
        result_str = generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, 
                                           scores_array, current_time, display_score_type)
        pymol_commands = generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, display_score_type)
        
        # Create chain-specific PDB with scores in B-factor
        scored_pdb = create_chain_specific_pdb(pdb_path, segment, selected_scores, protein_residues)
        
        # Create prediction file
        prediction_file = f"{pdb_id}_{display_score_type.lower()}_binding_site_residues.txt"
        with open(prediction_file, "w") as f:
            f.write(result_str)
        
        scored_pdb_name = f"{pdb_id}_{segment}_{display_score_type.lower()}_predictions_scores.pdb"
        os.rename(scored_pdb, scored_pdb_name)
        
        return mol_vis, pymol_commands, [prediction_file, scored_pdb_name]

    def fetch_interface(mode, pdb_id, pdb_file):
        if mode == "PDB ID":
            return fetch_pdb(pdb_id)
        elif mode == "Upload File":
            _, ext = os.path.splitext(pdb_file.name)
            file_path = os.path.join('./', f"{_}{ext}")
            if ext == '.cif':
                pdb_path = convert_cif_to_pdb(file_path)
            else:
                pdb_path= file_path
            return pdb_path

    def toggle_mode(selected_mode):
        if selected_mode == "PDB ID":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    
        
    mode.change(
        toggle_mode,
        inputs=[mode],
        outputs=[pdb_input, pdb_file]
    )

    prediction_btn.click(
        process_interface, 
        inputs=[mode, pdb_input, pdb_file, segment_input, score_type], 
        outputs=[predictions_output, molecule_output, download_output, 
                raw_scores_state, norm_scores_state, last_pdb_path, last_segment, last_pdb_id]
    )
    
    # Update visualization, PyMOL commands, and files when score type changes
    score_type.change(
        update_visualization_and_files,
        inputs=[score_type, raw_scores_state, norm_scores_state, last_pdb_path, last_segment, last_pdb_id],
        outputs=[molecule_output, predictions_output, download_output]
    )

    visualize_btn.click(
        fetch_interface, 
        inputs=[mode, pdb_input, pdb_file], 
        outputs=molecule_output2
    )

    gr.Markdown("## Examples")
    gr.Examples(
        examples=[
            ["7RPZ", "A"],
            ["2IWI", "B"],
            ["7LCJ", "R"],
            ["4OBE", "A"]
        ],
        inputs=[pdb_input, segment_input],
        outputs=[predictions_output, molecule_output, download_output]
    )

    def predict_utils(sequence):
        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()
    
        raw_scores = expit(outputs[:, 1] - outputs[:, 0])
        normalized_scores = normalize_scores(raw_scores)
    
        return {
        "raw_scores": raw_scores.tolist(),
        "normalized_scores": normalized_scores.tolist()
    }

    dummy_input = gr.Textbox(visible=False)
    dummy_output = gr.Textbox(visible=False)
    
    dummy_btn = gr.Button("Predict Sequence", visible=False)
    dummy_btn.click(
        predict_utils,
        inputs=[dummy_input],
        outputs=[dummy_output]
    )

demo.launch(share=True)