Update app.py

app.py

@@ -9,21 +9,23 @@ from transformers import BertTokenizer, BertModel
 from math import expm1
 
 # =====================
-# Load AMP Classifier
+# Load AMP Classifier Model (Random Forest)
 # =====================
+# Ensure 'RF.joblib' and 'norm (4).joblib' are in the same directory or provide full paths
 model = joblib.load("RF.joblib")
 scaler = joblib.load("norm (4).joblib")
 
 # =====================
-# Load ProtBert Globally
+# Load ProtBert Model Globally for MIC Prediction
 # =====================
 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()
+protbert_model = protbert_model.to(device).eval() # Set to evaluation mode
 
 # =====================
-# Feature List (ProPy)
+# Feature List (ProPy Descriptors) used by AMP Classifier
 # =====================
 selected_features = [
     "_SolventAccessibilityC3", "_SecondaryStrC1", "_SecondaryStrC3", "_ChargeC1", "_PolarityC1",
@@ -61,133 +63,222 @@ selected_features = [
 ]
 
 # =====================
-# AMP Feature Extractor
+# AMP Feature Extractor Function
 # =====================
 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."
+        return "Error: Sequence too short or invalid. Must contain at least 10 valid amino acids."
 
+    # Calculate various ProPy features
     dipeptide_features = AAComposition.CalculateAADipeptideComposition(sequence)
-    filtered_dipeptide_features = {k: dipeptide_features[k] for k in list(dipeptide_features.keys())[:420]}
+    # 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.
     ctd_features = CTD.CalculateCTD(sequence)
-    auto_features = Autocorrelation.CalculateAutoTotal(sequence)
-    pseudo_features = PseudoAAC.GetAPseudoAAC(sequence, lamda=9)
+    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
     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)
-    normalized_df = pd.DataFrame(normalized_array, columns=feature_df_all.columns)
-    selected_df = normalized_df[selected_features].fillna(0)
+    
+    # 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)
 
     return selected_df.values
 
 # =====================
-# 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 (ProtBert-based)
+# MIC Predictor Function (ProtBert-based)
 # =====================
-def predictmic(sequence):
+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.
+    """
     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."}
+        return {"Error": "Sequence too short or invalid for MIC prediction."}
 
-    # Tokenize
+    # Tokenize the sequence for ProtBert
     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)
 
-    # MIC model config
+    # Configuration for MIC models (paths to joblib files)
     bacteria_config = {
-        "E.coli": {
-            "model": "coli_xgboost_model.pkl",
-            "scaler": "coli_scaler.pkl",
-            "pca": None
+        "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
         },
-        "S.aureus": {
-            "model": "aur_xgboost_model.pkl",
-            "scaler": "aur_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
         },
-        "P.aeruginosa": {
-            "model": "arg_xgboost_model.pkl",
-            "scaler": "arg_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
         },
-        "K.Pneumonia": {
-            "model": "pne_mlp_model.pkl",
-            "scaler": "pne_scaler.pkl",
-            "pca": "pne_pca.pkl"
+        "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"
         }
     }
 
     mic_results = {}
-    for bacterium, cfg in bacteria_config.items():
+    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
+
         try:
-            scaler = joblib.load(cfg["scaler"])
-            scaled = scaler.transform(embedding)
-            if cfg["pca"]:
-                pca = joblib.load(cfg["pca"])
-                transformed = pca.transform(scaled)
+            # 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)
             else:
-                transformed = scaled
-            model = joblib.load(cfg["model"])
-            mic_log = model.predict(transformed)[0]
-            mic = round(expm1(mic_log), 3)
-            mic_results[bacterium] = mic
+                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
         except Exception as e:
-            mic_results[bacterium] = f"Error: {str(e)}"
+            mic_results[cfg["display_name"]] = f"Prediction Error: {str(e)}"
 
     return mic_results
 
 # =====================
-# Combined Prediction Function
+# Gradio Interface Functions
 # =====================
-def full_prediction(sequence):
+
+def amp_classifier_predict(sequence):
+    """
+    Function for AMP classification endpoint in Gradio.
+    Returns the AMP classification label, confidence, and SHAP plot Base64 string.
+    """
     features = extract_features(sequence)
-    if isinstance(features, str):
-        return "Error", "0%", {}
+    if isinstance(features, str): # Handle extraction error
+        return gr.Label(f"Error: {features}", label="AMP Classification"), None
+
     prediction = model.predict(features)[0]
     probabilities = model.predict_proba(features)[0]
-    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
+
+    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
 
 # =====================
-# Gradio Interface
+# Define Gradio Interface (hidden, for client connection)
 # =====================
-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)
+# 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)