app.py

import gradio as gr
import joblib
import numpy as np
import pandas as pd
from propy import AAComposition, CTD
import tensorflow as tf
from tensorflow.keras.models import load_model
import torch
from transformers import BertTokenizer, BertModel
from lime.lime_tabular import LimeTabularExplainer
from math import expm1

# Load AMP Classifier (Keras) and Scaler
model = load_model("Comb1_aac_ctd_RFE_selected_features_model.keras")
scaler = joblib.load("Comb1_aac_ctd_RFE_selected_features_scaler.joblib")

# Load ProtBert (for MIC prediction)
tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
protbert_model = BertModel.from_pretrained("Rostlab/prot_bert")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
protbert_model = protbert_model.to(device).eval()

# Define selected features (AAC + CTD, RFE-selected)
# Note: 'Activity' is the target label and is excluded from input features
selected_features = [
    '_PolarizabilityC1', '_PolarizabilityC2', '_PolarizabilityC3',
    '_SolventAccessibilityC1', '_SolventAccessibilityC2', '_SolventAccessibilityC3',
    '_SecondaryStrC1', '_SecondaryStrC2', '_SecondaryStrC3',
    '_ChargeC1', '_ChargeC2', '_ChargeC3',
    '_PolarityC1', '_PolarityC2', '_PolarityC3',
    '_NormalizedVDWVC1', '_NormalizedVDWVC2', '_NormalizedVDWVC3',
    '_HydrophobicityC1', '_HydrophobicityC2', '_HydrophobicityC3',
    '_PolarizabilityT12', '_PolarizabilityT13', '_PolarizabilityT23',
    '_SolventAccessibilityT12', '_SolventAccessibilityT13', '_SolventAccessibilityT23',
    '_SecondaryStrT12', '_SecondaryStrT13', '_SecondaryStrT23',
    '_ChargeT12', '_ChargeT13', '_ChargeT23',
    '_PolarityT12', '_PolarityT13', '_PolarityT23',
    '_NormalizedVDWVT12', '_NormalizedVDWVT13', '_NormalizedVDWVT23',
    '_HydrophobicityT12', '_HydrophobicityT13', '_HydrophobicityT23',
    'A', 'R', 'N', 'D', 'C', 'E', 'Q', 'G', 'H', 'I',
    'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V',
    'AA', 'AR', 'AN', 'AD', 'AC', 'AE', 'AQ', 'AG', 'AH', 'AI',
    'AL', 'AK', 'AM', 'AF', 'AP', 'AS', 'AT', 'AW', 'AY', 'AV',
    'RA', 'RR', 'RN', 'RD', 'RC', 'RE', 'RQ', 'RG', 'RH', 'RI',
    'RL', 'RK', 'RM', 'RF', 'RP', 'RS', 'RT', 'RW', 'RY', 'RV',
    'NA', 'NR', 'NN', 'ND', 'NC', 'NE', 'NQ', 'NG', 'NH', 'NI',
    'NL', 'NK', 'NM', 'NF', 'NP', 'NS', 'NT', 'NW', 'NY', 'NV',
    'DA', 'DR', 'DN', 'DD', 'DC', 'DE', 'DQ', 'DG', 'DH', 'DI',
    'DL', 'DK', 'DM', 'DF', 'DP', 'DS', 'DT', 'DW', 'DY', 'DV',
    'CA', 'CR', 'CN', 'CD', 'CC', 'CE', 'CQ', 'CG', 'CH', 'CI',
    'CL', 'CK', 'CM', 'CF', 'CP', 'CS', 'CT', 'CW', 'CY', 'CV',
    'EA', 'ER', 'EN', 'ED', 'EC', 'EE', 'EQ', 'EG', 'EH', 'EI',
    'EL', 'EK', 'EM', 'EF', 'EP', 'ES', 'ET', 'EW', 'EY', 'EV',
    'QA', 'QR', 'QN', 'QD', 'QC', 'QE', 'QQ', 'QG', 'QH', 'QI',
    'QL', 'QK', 'QM', 'QF', 'QP', 'QS', 'QT', 'QW', 'QY', 'QV',
    'GA', 'GR', 'GN', 'GD', 'GC', 'GE', 'GQ', 'GG', 'GH', 'GI',
    'GL', 'GK', 'GM', 'GF', 'GP', 'GS', 'GT', 'GW', 'GY', 'GV',
    'HA', 'HR', 'HN', 'HD', 'HC', 'HE', 'HQ', 'HG', 'HH', 'HI',
    'HL', 'HK', 'HM', 'HF', 'HP', 'HS', 'HT', 'HW', 'HY', 'HV',
    'IA', 'IR', 'IN', 'ID', 'IC', 'IE', 'IQ', 'IG', 'IH', 'II',
    'IL', 'IK', 'IM', 'IF', 'IP', 'IS', 'IT', 'IW', 'IY', 'IV',
    'LA', 'LR', 'LN', 'LD', 'LC', 'LE', 'LQ', 'LG', 'LH', 'LI',
    'LL', 'LK', 'LM', 'LF', 'LP', 'LS', 'LT', 'LW', 'LY', 'LV',
    'KA', 'KR', 'KN', 'KD', 'KC', 'KE', 'KQ', 'KG', 'KH', 'KI',
    'KL', 'KK', 'KM', 'KF', 'KP', 'KS', 'KT', 'KW', 'KY', 'KV',
    'MA', 'MR', 'MN', 'MD', 'MC', 'ME', 'MQ', 'MG', 'MH', 'MI',
    'ML', 'MK', 'MM', 'MF', 'MP', 'MS', 'MT', 'MW', 'MY', 'MV',
    'FA', 'FR', 'FN', 'FD', 'FC', 'FE', 'FQ', 'FG', 'FH', 'FI',
    'FL', 'FK', 'FM', 'FF', 'FP', 'FS', 'FT', 'FW', 'FY', 'FV',
    'PA', 'PR', 'PN', 'PD', 'PC', 'PE', 'PQ', 'PG', 'PH', 'PI',
    'PL', 'PK', 'PM', 'PF', 'PP', 'PS', 'PT', 'PW', 'PY', 'PV',
    'SA', 'SR', 'SN', 'SD', 'SC', 'SE', 'SQ', 'SG', 'SH', 'SI',
    'SL', 'SK', 'SM', 'SF', 'SP', 'SS', 'ST', 'SW', 'SY', 'SV',
    'TA', 'TR', 'TN', 'TD', 'TC', 'TE', 'TQ', 'TG', 'TH', 'TI',
    'TL', 'TK', 'TM', 'TF', 'TP', 'TS', 'TT', 'TW', 'TY', 'TV',
    'WA', 'WR', 'WN', 'WD', 'WC', 'WE', 'WQ', 'WG', 'WH', 'WI',
    'WL', 'WK', 'WM', 'WF', 'WP', 'WS', 'WT', 'WW', 'WY', 'WV',
    'YA', 'YR', 'YN', 'YD', 'YC', 'YE', 'YQ', 'YG', 'YH', 'YI',
    'YL', 'YK', 'YM', 'YF', 'YP', 'YS', 'YT', 'YW', 'YY', 'YV',
    'VA', 'VR', 'VN', 'VD', 'VC', 'VE', 'VQ', 'VG', 'VH', 'VI',
    'VL', 'VK', 'VM', 'VF', 'VP', 'VS', 'VT', 'VW', 'VY', 'VV'
]

# Wrapper to make Keras model behave like a sklearn classifier for LIME
def keras_predict_proba(X):
    """Return probabilities for both classes as [P(Non-AMP), P(AMP)]."""
    preds = model.predict(X, verbose=0)
    if preds.ndim == 1 or preds.shape[1] == 1:
        preds = preds.reshape(-1, 1)
        # Assuming sigmoid output = P(AMP); adjust if your model is reversed.
        return np.hstack([1 - preds, preds])
    return preds

# Dummy data for LIME
sample_data = np.random.rand(100, len(selected_features))
explainer = LimeTabularExplainer(
    training_data=sample_data,
    feature_names=selected_features,
    class_names=["Non-AMP", "AMP"],
    mode="classification"
)

# Feature extraction function (AAC + CTD only)
def extract_features(sequence):
    sequence = ''.join([aa for aa in sequence.upper() if aa in "ACDEFGHIKLMNPQRSTVWY"])
    if len(sequence) < 10:
        return "Error: Sequence too short."

    try:
        # AAC: 20 single AAs + 400 dipeptides = 420 features
        dipeptide_features = AAComposition.CalculateAADipeptideComposition(sequence)
        filtered_aac = {k: dipeptide_features[k] for k in list(dipeptide_features.keys())[:420]}

        # CTD: Composition, Transition, Distribution
        ctd_features = CTD.CalculateCTD(sequence)

        all_features_dict = {}
        all_features_dict.update(ctd_features)
        all_features_dict.update(filtered_aac)

        feature_df_all = pd.DataFrame([all_features_dict])
        normalized_array = scaler.transform(feature_df_all.values)
        normalized_df = pd.DataFrame(normalized_array, columns=feature_df_all.columns)

        if not set(selected_features).issubset(normalized_df.columns):
            missing = set(selected_features) - set(normalized_df.columns)
            return f"Error: Missing features: {list(missing)[:5]}..."

        selected_df = normalized_df[selected_features].fillna(0)
        return selected_df.values.astype(np.float32)
    except Exception as e:
        return f"Error in feature extraction: {str(e)}"

# MIC prediction function (unchanged)
def predictmic(sequence):
    sequence = ''.join([aa for aa in sequence.upper() if aa in "ACDEFGHIKLMNPQRSTVWY"])
    if len(sequence) < 10:
        return {"Error": "Sequence too short or invalid."}

    seq_spaced = ' '.join(list(sequence))
    tokens = tokenizer(seq_spaced, return_tensors="pt", padding='max_length', truncation=True, max_length=512)
    tokens = {k: v.to(device) for k, v in tokens.items()}

    with torch.no_grad():
        outputs = protbert_model(**tokens)
        embedding = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy().reshape(1, -1)

    bacteria_config = {
        "E.coli": {"model": "coli_xgboost_model.pkl", "scaler": "coli_scaler.pkl", "pca": None},
        "S.aureus": {"model": "aur_xgboost_model.pkl", "scaler": "aur_scaler.pkl", "pca": None},
        "P.aeruginosa": {"model": "arg_xgboost_model.pkl", "scaler": "arg_scaler.pkl", "pca": None},
        "K.Pneumonia": {"model": "pne_mlp_model.pkl", "scaler": "pne_scaler.pkl", "pca": "pne_pca.pkl"}
    }

    mic_results = {}
    for bacterium, cfg in bacteria_config.items():
        try:
            mic_scaler = joblib.load(cfg["scaler"])
            scaled = mic_scaler.transform(embedding)
            transformed = joblib.load(cfg["pca"]).transform(scaled) if cfg["pca"] else scaled
            mic_model = joblib.load(cfg["model"])
            mic_log = mic_model.predict(transformed)[0]
            mic = round(expm1(mic_log), 3)
            mic_results[bacterium] = mic
        except Exception as e:
            mic_results[bacterium] = f"Error: {str(e)}"

    return mic_results

# Main prediction function
def full_prediction(sequence):
    features = extract_features(sequence)
    if isinstance(features, str):
        return features

    # Keras prediction
    raw_pred = model.predict(features, verbose=0)

    # Handle sigmoid (1 output) vs softmax (>=2 outputs)
    if raw_pred.ndim == 1 or raw_pred.shape[1] == 1:
        prob_amp = float(raw_pred.flatten()[0])  # assume output = P(AMP)
        if prob_amp >= 0.5:
            prediction = 1  # AMP
            confidence = round(prob_amp * 100, 2)
        else:
            prediction = 0  # Non-AMP
            confidence = round((1 - prob_amp) * 100, 2)
    else:
        class_idx = int(np.argmax(raw_pred[0]))
        prediction = class_idx
        confidence = round(float(raw_pred[0][class_idx]) * 100, 2)

    # Label convention: 1 = AMP, 0 = Non-AMP (swap if your model uses the opposite)
    amp_result = "Antimicrobial Peptide (AMP)" if prediction == 1 else "Non-AMP"
    result = f"Prediction: {amp_result}\nConfidence: {confidence}%\n"

    if prediction == 1:
        mic_values = predictmic(sequence)
        result += "\nPredicted MIC Values (μM):\n"
        for org, mic in mic_values.items():
            result += f"- {org}: {mic}\n"
    else:
        result += "\nMIC prediction skipped for Non-AMP sequences.\n"

    # LIME explanation
    try:
        explanation = explainer.explain_instance(
            data_row=features[0],
            predict_fn=keras_predict_proba,
            num_features=10
        )

        result += "\nTop Features Influencing Prediction:\n"
        for feat, weight in explanation.as_list():
            result += f"- {feat}: {round(weight, 4)}\n"
    except Exception as e:
        result += f"\nLIME explanation failed: {str(e)}\n"

    return result

# Gradio UI
iface = gr.Interface(
    fn=full_prediction,
    inputs=gr.Textbox(label="Enter Protein Sequence"),
    outputs=gr.Textbox(label="Results"),
    title="AMP & MIC Predictor + LIME Explanation",
    description="Paste an amino acid sequence (≥10 characters). Get AMP classification, MIC predictions, and LIME interpretability insights."
)

iface.launch(share=True)