#!/usr/bin/env python3 """ Epitope Prediction Model Interface Loads and uses the trained deep learning model for epitope prediction """ import os import logging # Handle TensorFlow import gracefully for deployment try: import tensorflow as tf from tensorflow import keras TF_AVAILABLE = True logger = logging.getLogger(__name__) logger.info(f"TensorFlow {tf.__version__} loaded successfully") except ImportError as e: TF_AVAILABLE = False tf = None keras = None logger = logging.getLogger(__name__) logger.warning(f"TensorFlow not available: {e}") try: import numpy as np except ImportError: logger.error("NumPy is required but not available") raise # Optional imports for deployment compatibility try: import sklearn SKLEARN_AVAILABLE = True except ImportError: SKLEARN_AVAILABLE = False logging.warning("scikit-learn not available - some features may be limited") # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class EpitopePredictor: """ Epitope prediction class that loads the trained model and performs predictions """ def __init__(self, model_path=None): """ Initialize the predictor with the trained model Args: model_path (str): Path to the model file. If None, uses default path. """ # Check if TensorFlow is available if not TF_AVAILABLE: logger.error("TensorFlow is not available. Model prediction will not work.") self.model = None return # Amino acid to number mapping (same as used in training) self.amino_acid_to_num = { 'A': 1, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6, 'H': 7, 'I': 8, 'K': 9, 'L': 10, 'M': 11, 'N': 12, 'P': 13, 'Q': 14, 'R': 15, 'S': 16, 'T': 17, 'V': 18, 'W': 19, 'Y': 20, 'X': 0 # X for unknown } # Class mapping (same as used in training) self.class_mapping = { 'B_cell_negative': 0, 'B_cell_positive': 1, 'T_cell_MHC_negative': 2, 'T_cell_MHC_positive': 3 } # Reverse mapping for predictions self.idx_to_class = {v: k for k, v in self.class_mapping.items()} # Model parameters self.window_size = 20 self.step_size = 1 self.confidence_threshold = 0.5 # Load the model try: self.model = self._load_model(model_path) except Exception as e: logger.error(f"Failed to load model: {e}") self.model = None def _load_model(self, model_path=None): """ Load the trained model Args: model_path (str): Path to model file Returns: tensorflow.keras.Model: Loaded model """ if model_path is None: # Try different model formats in order of preference possible_paths = [ 'models/epitope_model.keras', 'models/epitope_model.h5', 'models/epitope_model_savedmodel', '../epitope_model.keras', '../epitope_model.h5', '../epitope_model_savedmodel', 'epitope_model.keras', 'epitope_model.h5', 'epitope_model_savedmodel' ] for path in possible_paths: if os.path.exists(path): model_path = path break if model_path is None: raise FileNotFoundError("No trained model found. Please ensure the model file exists.") try: logger.info(f"Loading model from: {model_path}") if model_path.endswith('.keras') or model_path.endswith('.h5'): model = keras.models.load_model(model_path, compile=False) else: # Assume SavedModel format model = tf.saved_model.load(model_path) logger.info("Model loaded successfully") # Test the model with a dummy input to ensure it works try: dummy_input = np.array([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]]) if hasattr(model, 'predict'): _ = model.predict(dummy_input, verbose=0) else: _ = model(dummy_input) logger.info("Model test prediction successful") except Exception as test_error: logger.warning(f"Model test failed: {test_error}") return model except Exception as e: logger.error(f"Error loading model: {e}") raise RuntimeError(f"Failed to load model from {model_path}: {e}") def encode_sequence(self, sequence): """ Encode amino acid sequence to numerical representation Args: sequence (str): Amino acid sequence Returns: list: Encoded sequence """ return [self.amino_acid_to_num.get(aa.upper(), 0) for aa in sequence] def sliding_window_prediction(self, sequence): """ Perform sliding window prediction on a protein sequence Args: sequence (str): Protein sequence Returns: tuple: (b_cell_epitopes, t_cell_epitopes) lists with (epitope, confidence, position_range) """ logger.debug(f"Starting sliding window prediction for sequence of length {len(sequence)}") b_cell_epitopes = [] t_cell_epitopes = [] if len(sequence) < self.window_size: logger.warning(f"Sequence too short ({len(sequence)} < {self.window_size})") return b_cell_epitopes, t_cell_epitopes # Prepare batch data for efficient prediction subseq_list = [] positions = [] for i in range(0, len(sequence) - self.window_size + 1, self.step_size): sub_seq = sequence[i:i + self.window_size] encoded_sub_seq = self.encode_sequence(sub_seq) subseq_list.append(encoded_sub_seq) positions.append((i, i + self.window_size)) if not subseq_list: logger.warning("No subsequences generated for prediction") return b_cell_epitopes, t_cell_epitopes logger.debug(f"Generated {len(subseq_list)} subsequences for prediction") # Convert to numpy array for batch prediction padded_subseq_array = np.array(subseq_list) try: logger.debug(f"Performing batch prediction on {padded_subseq_array.shape} array") # Perform batch predictions if hasattr(self.model, 'predict'): predicted_probs = self.model.predict(padded_subseq_array, batch_size=64, verbose=0) else: # For SavedModel format predicted_probs = self.model(padded_subseq_array).numpy() logger.debug(f"Model prediction completed, output shape: {predicted_probs.shape}") # Process predictions for i, (probs, (start_pos, end_pos)) in enumerate(zip(predicted_probs, positions)): predicted_class = np.argmax(probs) confidence = np.max(probs) predicted_label = self.idx_to_class[predicted_class] # Only include predictions above threshold if confidence >= self.confidence_threshold: sub_seq = sequence[start_pos:end_pos] pos_range = f"{start_pos+1}-{end_pos}" # 1-based indexing for display if predicted_label == "B_cell_positive": b_cell_epitopes.append((sub_seq, float(confidence), pos_range)) elif predicted_label == "T_cell_MHC_positive": t_cell_epitopes.append((sub_seq, float(confidence), pos_range)) logger.debug(f"Prediction processing completed: {len(b_cell_epitopes)} B-cell, {len(t_cell_epitopes)} T-cell epitopes above threshold") except Exception as e: logger.error(f"Error during prediction: {e}") raise return b_cell_epitopes, t_cell_epitopes def predict_epitopes(self, sequence, threshold=None): """ Main prediction function Args: sequence (str): Protein sequence threshold (float): Confidence threshold (optional) Returns: tuple: (b_cell_epitopes, t_cell_epitopes) """ # Check if model is available if self.model is None: logger.warning("Model not available, returning demo predictions") return self._generate_demo_predictions(sequence) logger.info(f"Starting epitope prediction for sequence of length {len(sequence)}") if threshold is not None: original_threshold = self.confidence_threshold self.confidence_threshold = threshold logger.debug(f"Using custom threshold: {threshold}") try: # Clean the sequence clean_sequence = ''.join(c.upper() for c in sequence if c.upper() in self.amino_acid_to_num) if len(clean_sequence) != len(sequence): logger.warning(f"Sequence contained invalid characters. Cleaned: {len(sequence)} -> {len(clean_sequence)}") # Perform prediction b_cell_epitopes, t_cell_epitopes = self.sliding_window_prediction(clean_sequence) # Sort by confidence (highest first) b_cell_epitopes.sort(key=lambda x: x[1], reverse=True) t_cell_epitopes.sort(key=lambda x: x[1], reverse=True) logger.info(f"Prediction completed: {len(b_cell_epitopes)} B-cell, {len(t_cell_epitopes)} T-cell epitopes found") return b_cell_epitopes, t_cell_epitopes finally: if threshold is not None: self.confidence_threshold = original_threshold def _generate_demo_predictions(self, sequence): """ Generate demo predictions when model is not available Args: sequence (str): Protein sequence Returns: tuple: (b_cell_epitopes, t_cell_epitopes) with demo data """ import random random.seed(42) # For consistent demo results b_cell_epitopes = [] t_cell_epitopes = [] # Generate some demo epitopes seq_len = len(sequence) if seq_len >= 20: # Generate a few demo B-cell epitopes for i in range(0, min(seq_len - 19, 3)): start = i * 25 if start + 20 <= seq_len: epitope = sequence[start:start + 20] confidence = 0.6 + random.random() * 0.3 # 0.6-0.9 pos_range = f"{start + 1}-{start + 20}" b_cell_epitopes.append((epitope, confidence, pos_range)) # Generate a few demo T-cell epitopes for i in range(1, min(seq_len - 19, 3)): start = i * 30 + 10 if start + 20 <= seq_len: epitope = sequence[start:start + 20] confidence = 0.5 + random.random() * 0.4 # 0.5-0.9 pos_range = f"{start + 1}-{start + 20}" t_cell_epitopes.append((epitope, confidence, pos_range)) logger.info(f"Demo predictions generated: {len(b_cell_epitopes)} B-cell, {len(t_cell_epitopes)} T-cell epitopes") return b_cell_epitopes, t_cell_epitopes def get_sequence_markup(self, sequence, epitopes, epitope_type='B-cell'): """ Generate sequence markup for visualization Args: sequence (str): Original sequence epitopes (list): List of epitopes with positions epitope_type (str): Type of epitopes ('B-cell' or 'T-cell') Returns: str: Marked up sequence """ markup = ['.' for _ in sequence] # Default to non-epitope for epitope, confidence, pos_range in epitopes: start, end = map(int, pos_range.split('-')) start -= 1 # Convert to 0-based indexing end -= 1 # Mark epitope positions marker = 'E' if epitope_type == 'B-cell' else 'T' for i in range(start, min(end + 1, len(markup))): markup[i] = marker return ''.join(markup) def set_confidence_threshold(self, threshold): """ Set the confidence threshold for predictions Args: threshold (float): New threshold value (0.0 to 1.0) """ if 0.0 <= threshold <= 1.0: self.confidence_threshold = threshold else: raise ValueError("Threshold must be between 0.0 and 1.0") def get_model_info(self): """ Get information about the loaded model Returns: dict: Model information """ info = { 'window_size': self.window_size, 'step_size': self.step_size, 'confidence_threshold': self.confidence_threshold, 'classes': list(self.class_mapping.keys()), 'amino_acids': list(self.amino_acid_to_num.keys()) } if hasattr(self.model, 'summary'): try: # Get model summary as string import io import sys old_stdout = sys.stdout sys.stdout = buffer = io.StringIO() self.model.summary() sys.stdout = old_stdout info['model_summary'] = buffer.getvalue() except: info['model_summary'] = "Model summary not available" return info