Update app.py

app.py

@@ -8,25 +8,17 @@ import torch
 from transformers import BertTokenizer, BertModel
 from math import expm1
 
-# =====================
-# Load AMP Classifier Model (Random Forest)
-# =====================
-# Ensure 'RF.joblib' and 'norm (4).joblib' are in the same directory or provide full paths
+# Load AMP Classifier
 model = joblib.load("RF.joblib")
 scaler = joblib.load("norm (4).joblib")
 
-# =====================
-# Load ProtBert Model Globally for MIC Prediction
-# =====================
+# Load ProtBert Globally
 tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
 protbert_model = BertModel.from_pretrained("Rostlab/prot_bert")
-# Move model to GPU if available for faster inference
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-protbert_model = protbert_model.to(device).eval() # Set to evaluation mode
+protbert_model = protbert_model.to(device).eval()
 
-# =====================
-# Feature List (ProPy Descriptors) used by AMP Classifier
-# =====================
+# Selected Features
 selected_features = [
     "_SolventAccessibilityC3", "_SecondaryStrC1", "_SecondaryStrC3", "_ChargeC1", "_PolarityC1",
     "_NormalizedVDWVC1", "_HydrophobicityC3", "_SecondaryStrT23", "_PolarizabilityD1001",
@@ -62,223 +54,113 @@ selected_features = [
     "APAAC24"
 ]
 
-# =====================
-# AMP Feature Extractor Function
-# =====================
+# AMP Feature Extractor
 def extract_features(sequence):
-    """
-    Extracts physiochemical and compositional features from a protein sequence using ProPy.
-    Applies the pre-trained scaler and selects relevant features.
-    """
     all_features_dict = {}
-    # Clean sequence to include only valid amino acids
     sequence = ''.join([aa for aa in sequence.upper() if aa in "ACDEFGHIKLMNPQRSTVWY"])
-
     if len(sequence) < 10:
-        return "Error: Sequence too short or invalid. Must contain at least 10 valid amino acids."
-
-    # Calculate various ProPy features
+        return "Error: Sequence too short."
     dipeptide_features = AAComposition.CalculateAADipeptideComposition(sequence)
-    # Note: Dipeptide composition calculates 400 features, using a slice here might be specific to the original model's training
-    # If the original model used all 400, this slice needs to be adjusted or removed.
-    # For now, keeping as per the provided code.
-    filtered_dipeptide_features = {k: dipeptide_features[k] for k in list(dipeptide_features.keys())[:420]} # This slice is unusual if only 400 dipeptides exist.
+    filtered_dipeptide_features = {k: dipeptide_features[k] for k in list(dipeptide_features.keys())[:420]}
     ctd_features = CTD.CalculateCTD(sequence)
-    auto_features = Autocorrelation.CalculateAutoTotal(sequence) # Includes Moran, Geary, Moreau-Broto
-    pseudo_features = PseudoAAC.GetAPseudoAAC(sequence, lamda=9) # Pseudo Amino Acid Composition
-
-    # Combine all extracted features into a single dictionary
+    auto_features = Autocorrelation.CalculateAutoTotal(sequence)
+    pseudo_features = PseudoAAC.GetAPseudoAAC(sequence, lamda=9)
     all_features_dict.update(ctd_features)
     all_features_dict.update(filtered_dipeptide_features)
     all_features_dict.update(auto_features)
     all_features_dict.update(pseudo_features)
-
-    # Convert to DataFrame for consistent column handling with scaler
     feature_df_all = pd.DataFrame([all_features_dict])
-
-    # Handle missing features (if any arise from short sequences or specific AA combinations not producing all features)
-    # Ensure all selected_features are present, add as 0 if missing.
-    for col in selected_features:
-        if col not in feature_df_all.columns:
-            feature_df_all[col] = 0
-
-    # Normalize features using the pre-trained scaler
-    # Ensure the order of columns matches the scaler's training order before scaling
-    feature_df_all = feature_df_all[scaler.feature_names_in_] # Align columns with scaler's expected input
     normalized_array = scaler.transform(feature_df_all.values)
-    
-    # Select only the features that the final RF model expects
-    selected_df = pd.DataFrame(normalized_array, columns=scaler.feature_names_in_)[selected_features].fillna(0)
-
+    normalized_df = pd.DataFrame(normalized_array, columns=feature_df_all.columns)
+    selected_df = normalized_df[selected_features].fillna(0)
     return selected_df.values
 
-# =====================
-# MIC Predictor Function (ProtBert-based)
-# =====================
-def predict_mic_values(sequence, selected_bacteria_keys):
-    """
-    Predicts Minimum Inhibitory Concentration (MIC) for a given peptide sequence
-    against selected bacteria using ProtBert embeddings and pre-trained models.
-    """
+# AMP Classifier
+def predict(sequence):
+    features = extract_features(sequence)
+    if isinstance(features, str):
+        return features
+    prediction = model.predict(features)[0]
+    probabilities = model.predict_proba(features)[0]
+    if prediction == 0:
+        return f"{probabilities[0] * 100:.2f}% chance of being an Antimicrobial Peptide (AMP)"
+    else:
+        return f"{probabilities[1] * 100:.2f}% chance of being Non-AMP"
+
+# MIC Predictor
+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 for MIC prediction."}
-
-    # Tokenize the sequence for ProtBert
+        return {"Error": "Sequence too short or invalid. Must contain at least 10 valid amino acids."}
     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()}
-
-    # Get ProtBert embedding
     with torch.no_grad():
         outputs = protbert_model(**tokens)
-        # Use mean of last hidden state as sequence embedding
         embedding = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy().reshape(1, -1)
-
-    # Configuration for MIC models (paths to joblib files)
     bacteria_config = {
-        "e_coli": { # Changed keys to match frontend values
-            "display_name": "E.coli",
-            "model_path": "coli_xgboost_model.pkl",
-            "scaler_path": "coli_scaler.pkl",
-            "pca_path": None
+        "E.coli": {
+            "model": "coli_xgboost_model.pkl",
+            "scaler": "coli_scaler.pkl",
+            "pca": None
         },
-        "s_aureus": { # Changed keys to match frontend values
-            "display_name": "S.aureus",
-            "model_path": "aur_xgboost_model.pkl",
-            "scaler_path": "aur_scaler.pkl",
-            "pca_path": None
+        "S.aureus": {
+            "model": "aur_xgboost_model.pkl",
+            "scaler": "aur_scaler.pkl",
+            "pca": None
         },
-        "p_aeruginosa": { # Changed keys to match frontend values
-            "display_name": "P.aeruginosa",
-            "model_path": "arg_xgboost_model.pkl",
-            "scaler_path": "arg_scaler.pkl",
-            "pca_path": None
+        "P.aeruginosa": {
+            "model": "arg_xgboost_model.pkl",
+            "scaler": "arg_scaler.pkl",
+            "pca": None
         },
-        "k_pneumoniae": { # Changed keys to match frontend values
-            "display_name": "K.Pneumoniae",
-            "model_path": "pne_mlp_model.pkl",
-            "scaler_path": "pne_scaler.pkl",
-            "pca_path": "pne_pca.pkl"
+        "K.Pneumonia": {
+            "model": "pne_mlp_model.pkl",
+            "scaler": "pne_scaler.pkl",
+            "pca": "pne_pca.pkl"
         }
     }
-
     mic_results = {}
-    for bacterium_key in selected_bacteria_keys:
-        cfg = bacteria_config.get(bacterium_key)
-        if not cfg:
-            mic_results[bacterium_key] = "Error: Invalid bacterium key"
-            continue
-
+    for bacterium, cfg in bacteria_config.items():
         try:
-            # Load scaler and transform embedding
-            scaler = joblib.load(cfg["scaler_path"])
-            scaled_embedding = scaler.transform(embedding)
-
-            # Apply PCA if configured
-            if cfg["pca_path"]:
-                pca = joblib.load(cfg["pca_path"])
-                final_features = pca.transform(scaled_embedding)
+            scaler = joblib.load(cfg["scaler"])
+            scaled = scaler.transform(embedding)
+            if cfg["pca"]:
+                pca = joblib.load(cfg["pca"])
+                transformed = pca.transform(scaled)
             else:
-                final_features = scaled_embedding
-
-            # Load and predict with the MIC model
-            mic_model = joblib.load(cfg["model_path"])
-            mic_log = mic_model.predict(final_features)[0]
-
-            # Convert log-transformed MIC back to original scale (µM)
-            mic = round(expm1(mic_log), 3) # expm1(x) is equivalent to exp(x) - 1, robust for small x
-            mic_results[cfg["display_name"]] = mic
+                transformed = scaled
+            model = joblib.load(cfg["model"])
+            mic_log = model.predict(transformed)[0]
+            mic = round(expm1(mic_log), 3)
+            mic_results[bacterium] = mic
         except Exception as e:
-            mic_results[cfg["display_name"]] = f"Prediction Error: {str(e)}"
-
+            mic_results[bacterium] = f"Error: {str(e)}"
     return mic_results
 
-# =====================
-# Gradio Interface Functions
-# =====================
-
-def amp_classifier_predict(sequence):
-    """
-    Function for AMP classification endpoint in Gradio.
-    Returns the AMP classification label, confidence, and SHAP plot Base64 string.
-    """
+# Combined Prediction
+def full_prediction(sequence):
     features = extract_features(sequence)
-    if isinstance(features, str): # Handle extraction error
-        return gr.Label(f"Error: {features}", label="AMP Classification"), None
-
+    if isinstance(features, str):
+        return "Error", "0%", {}
     prediction = model.predict(features)[0]
     probabilities = model.predict_proba(features)[0]
-
-    amp_label = "Antimicrobial Peptide (AMP)" if prediction == 0 else "Non-AMP"
-    confidence_value = probabilities[prediction] # Confidence of the predicted class
-
-    # Placeholder for SHAP plot generation (not implemented in this snippet)
-    # In a real scenario, you'd generate a SHAP plot image here (e.g., using matplotlib, shap library)
-    # and encode it to base64.
-    shap_plot_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" # A tiny transparent PNG base64
-
-    # The Gradio `predict` function can return structured data as a dictionary if using `gr.JSON` output
-    # However, since the frontend is expecting `data[0].label`, `data[0].confidence`, etc.
-    # we'll return a dictionary that matches that structure.
-    return {
-        "label": amp_label,
-        "confidence": confidence_value,
-        "shap_plot_base64": shap_plot_base64 # Return SHAP plot as Base64 (placeholder for now)
-    }
-
-def mic_predictor_predict(sequence, selected_bacteria):
-    """
-    Function for MIC prediction endpoint in Gradio.
-    Takes the sequence and a list of selected bacteria keys.
-    """
-    # Only predict MIC if AMP (Positive) classification
-    # This check would ideally be part of the frontend logic or a combined backend function
-    # but for standalone MIC endpoint, we just proceed.
-    # The frontend is responsible for calling this only if AMP is positive.
-    mic_results = predict_mic_values(sequence, selected_bacteria)
-    return mic_results # Returns a dictionary of MIC values
-
-# =====================
-# Define Gradio Interface (hidden, for client connection)
-# =====================
-# This Gradio app is designed to be used as a backend service by your custom HTML frontend.
-# The inputs and outputs here correspond to what the frontend's `gradio.client` expects.
-
-with gr.Blocks() as demo:
-    gr.Markdown("# BCBU-ZC AMP/MIC Backend Service")
-    gr.Markdown("This Gradio application serves as the backend for the AMP classification and MIC prediction. It provides endpoints for sequence analysis and MIC prediction.")
-
-    with gr.Tab("AMP Classification"):
-        gr.Markdown("### AMP Classification Endpoint (`/predict`)")
-        amp_input_sequence = gr.Textbox(label="Amino Acid Sequence")
-        amp_output_json = gr.JSON(label="Classification Result (Label, Confidence, SHAP Plot Base64)")
-        amp_predict_button = gr.Button("Predict AMP")
-        amp_predict_button.click(
-            fn=amp_classifier_predict,
-            inputs=[amp_input_sequence],
-            outputs=[amp_output_json],
-            api_name="predict" # Define an API endpoint name for `gradio.client`
-        )
-
-    with gr.Tab("MIC Prediction"):
-        gr.Markdown("### MIC Prediction Endpoint (`/predict_mic`)")
-        mic_input_sequence = gr.Textbox(label="Amino Acid Sequence")
-        mic_selected_bacteria = gr.CheckboxGroup(
-            label="Select Bacteria",
-            choices=["e_coli", "p_aeruginosa", "s_aureus", "k_pneumoniae"],
-            value=["e_coli", "p_aeruginosa", "s_aureus", "k_pneumoniae"] # Default for testing
-        )
-        mic_output_json = gr.JSON(label="Predicted MIC Values (µM)")
-        mic_predict_button = gr.Button("Predict MIC")
-        mic_predict_button.click(
-            fn=mic_predictor_predict,
-            inputs=[mic_input_sequence, mic_selected_bacteria],
-            outputs=[mic_output_json],
-            api_name="predict_mic" # Define a separate API endpoint name
-        )
-
-# Launch the Gradio app
-# `share=True` creates a public, temporary URL for external access (useful for testing frontend)
-# `allowed_paths` should be set to allow access from specific origins if deploying
-demo.launch(share=True)
+    amp_result = "Antimicrobial Peptide (AMP)" if prediction == 0 else "Non-AMP"
+    confidence = round(probabilities[0 if prediction == 0 else 1] * 100, 2)
+    mic_values = predictmic(sequence)
+    return amp_result, f"{confidence}%", mic_values
+
+# Gradio Interface
+iface = gr.Interface(
+    fn=full_prediction,
+    inputs=gr.Textbox(label="Enter Protein Sequence"),
+    outputs=[
+        gr.Label(label="AMP Classification"),
+        gr.Label(label="Confidence"),
+        gr.JSON(label="Predicted MIC (µM) for Each Bacterium")
+    ],
+    title="AMP & MIC Predictor",
+    description="Enter an amino acid sequence (≥10 valid letters) to predict AMP class and MIC values."
+)
+
+iface.launch(share=True)