Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import joblib | |
| import numpy as np | |
| from itertools import product | |
| import torch.nn as nn | |
| import shap | |
| import matplotlib.pyplot as plt | |
| import io | |
| import json | |
| from PIL import Image | |
| class VirusClassifier(nn.Module): | |
| def __init__(self, input_shape: int): | |
| super(VirusClassifier, self).__init__() | |
| self.network = nn.Sequential( | |
| nn.Linear(input_shape, 64), | |
| nn.GELU(), | |
| nn.BatchNorm1d(64), | |
| nn.Dropout(0.3), | |
| nn.Linear(64, 32), | |
| nn.GELU(), | |
| nn.BatchNorm1d(32), | |
| nn.Dropout(0.3), | |
| nn.Linear(32, 32), | |
| nn.GELU(), | |
| nn.Linear(32, 2) | |
| ) | |
| def forward(self, x): | |
| return self.network(x) | |
| def get_feature_importance(self, x): | |
| """Calculate feature importance using gradient-based method for the human class (index 1)""" | |
| x.requires_grad_(True) | |
| output = self.network(x) | |
| probs = torch.softmax(output, dim=1) | |
| # We focus on the human class (index 1) probability | |
| human_prob = probs[..., 1] | |
| human_prob.backward() | |
| # The gradient shows how each feature affects the human probability | |
| importance = x.grad | |
| return importance, float(human_prob) | |
| def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray: | |
| """Convert sequence to k-mer frequency vector""" | |
| kmers = [''.join(p) for p in product("ACGT", repeat=k)] | |
| kmer_dict = {km: i for i, km in enumerate(kmers)} | |
| vec = np.zeros(len(kmers), dtype=np.float32) | |
| for i in range(len(sequence) - k + 1): | |
| kmer = sequence[i:i+k] | |
| if kmer in kmer_dict: | |
| vec[kmer_dict[kmer]] += 1 | |
| total_kmers = len(sequence) - k + 1 | |
| if total_kmers > 0: | |
| vec = vec / total_kmers | |
| return vec | |
| def parse_fasta(text): | |
| sequences = [] | |
| current_header = None | |
| current_sequence = [] | |
| for line in text.split('\n'): | |
| line = line.strip() | |
| if not line: | |
| continue | |
| if line.startswith('>'): | |
| if current_header: | |
| sequences.append((current_header, ''.join(current_sequence))) | |
| current_header = line[1:] | |
| current_sequence = [] | |
| else: | |
| current_sequence.append(line.upper()) | |
| if current_header: | |
| sequences.append((current_header, ''.join(current_sequence))) | |
| return sequences | |
| def predict(file_obj): | |
| if file_obj is None: | |
| return "Please upload a FASTA file", None | |
| try: | |
| if isinstance(file_obj, str): | |
| text = file_obj | |
| else: | |
| text = file_obj.decode('utf-8') | |
| except Exception as e: | |
| return f"Error reading file: {str(e)}", None | |
| k = 4 | |
| kmers = [''.join(p) for p in product("ACGT", repeat=k)] | |
| kmer_dict = {km: i for i, km in enumerate(kmers)} | |
| try: | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| model = VirusClassifier(256).to(device) | |
| state_dict = torch.load('model.pt', map_location=device) | |
| model.load_state_dict(state_dict) | |
| scaler = joblib.load('scaler.pkl') | |
| model.eval() | |
| except Exception as e: | |
| return f"Error loading model: {str(e)}", None | |
| results_text = "" | |
| plot_image = None | |
| try: | |
| sequences = parse_fasta(text) | |
| header, seq = sequences[0] | |
| raw_freq_vector = sequence_to_kmer_vector(seq) | |
| kmer_vector = scaler.transform(raw_freq_vector.reshape(1, -1)) | |
| X_tensor = torch.FloatTensor(kmer_vector).to(device) | |
| # Calculate final probabilities first | |
| with torch.no_grad(): | |
| output = model(X_tensor) | |
| probs = torch.softmax(output, dim=1) | |
| human_prob = float(probs[0][1]) | |
| # Get feature importance using integrated gradients | |
| baseline = torch.zeros_like(X_tensor) # baseline of zeros | |
| steps = 50 | |
| all_importance = [] | |
| for i in range(steps + 1): | |
| alpha = i / steps | |
| interpolated = baseline + alpha * (X_tensor - baseline) | |
| interpolated.requires_grad_(True) | |
| output = model(interpolated) | |
| probs = torch.softmax(output, dim=1) | |
| human_class = probs[..., 1] | |
| if interpolated.grad is not None: | |
| interpolated.grad.zero_() | |
| human_class.backward() | |
| all_importance.append(interpolated.grad.cpu().numpy()) | |
| # Average the gradients | |
| kmer_importance = np.mean(all_importance, axis=0)[0] | |
| # Scale to match probability difference | |
| target_diff = human_prob - 0.5 # difference from neutral prediction | |
| current_sum = np.sum(kmer_importance) | |
| if current_sum != 0: # avoid division by zero | |
| kmer_importance = kmer_importance * (target_diff / current_sum) | |
| # Get top k-mers by absolute importance | |
| top_k = 10 | |
| top_indices = np.argsort(np.abs(kmer_importance))[-top_k:][::-1] | |
| important_kmers = [ | |
| { | |
| 'kmer': list(kmer_dict.keys())[list(kmer_dict.values()).index(i)], | |
| 'importance': float(kmer_importance[i]), | |
| 'frequency': float(raw_freq_vector[i]), | |
| 'scaled': float(kmer_vector[0][i]) | |
| } | |
| for i in top_indices | |
| ] | |
| # Prepare data for SHAP waterfall plot | |
| top_features = [item['kmer'] for item in important_kmers] | |
| top_values = [item['importance'] for item in important_kmers] | |
| # Calculate the impact of remaining features | |
| others_mask = np.ones_like(kmer_importance, dtype=bool) | |
| others_mask[top_indices] = False | |
| others_sum = np.sum(kmer_importance[others_mask]) | |
| top_features.append("Others") | |
| top_values.append(others_sum) | |
| # Calculate final probabilities first | |
| with torch.no_grad(): | |
| output = model(X_tensor) | |
| probs = torch.softmax(output, dim=1) | |
| human_prob = float(probs[0][1]) | |
| # Create SHAP explanation | |
| # We'll use the actual probabilities for alignment | |
| explanation = shap.Explanation( | |
| values=np.array(top_values), | |
| base_values=0.5, # Start from neutral prediction | |
| data=np.array([ | |
| raw_freq_vector[kmer_dict[feat]] if feat != "Others" | |
| else np.sum(raw_freq_vector[others_mask]) | |
| for feat in top_features | |
| ]), | |
| feature_names=top_features | |
| ) | |
| explanation.expected_value = 0.5 # Start from neutral prediction | |
| # Calculate step-by-step probabilities | |
| current_prob = 0.5 # Start at neutral | |
| steps = [('Start', current_prob, 0)] | |
| # Process each k-mer contribution | |
| for kmer in important_kmers: | |
| change = kmer['importance'] | |
| current_prob += change | |
| steps.append((kmer['kmer'], current_prob, change)) | |
| # Add final "Others" contribution | |
| current_prob += others_sum | |
| steps.append(('Others', current_prob, others_sum)) | |
| # Create step plot | |
| plt.figure(figsize=(12, 6)) | |
| x = range(len(steps)) | |
| y = [step[1] for step in steps] | |
| # Plot steps | |
| plt.step(x, y, 'b-', where='post', label='Probability', linewidth=2) | |
| plt.plot(x, y, 'b.', markersize=10) | |
| # Add reference line | |
| plt.axhline(y=0.5, color='r', linestyle='--', label='Neutral (0.5)') | |
| # Customize plot | |
| plt.grid(True, linestyle='--', alpha=0.7) | |
| plt.ylim(0, 1) | |
| plt.ylabel('Human Probability') | |
| plt.title(f'K-mer Contributions to Prediction (final prob: {human_prob:.3f})') | |
| # Add labels for each point | |
| for i, (kmer, prob, change) in enumerate(steps): | |
| # Add k-mer label | |
| plt.annotate(kmer, | |
| (i, prob), | |
| xytext=(0, 10 if i % 2 == 0 else -20), # Alternate up/down | |
| textcoords='offset points', | |
| ha='center', | |
| rotation=45 if len(kmer) > 5 else 0) | |
| # Add change value | |
| if i > 0: # Skip first point (Start) | |
| change_text = f'{change:+.3f}' | |
| color = 'green' if change > 0 else 'red' | |
| plt.annotate(change_text, | |
| (i, prob), | |
| xytext=(0, -20 if i % 2 == 0 else 10), | |
| textcoords='offset points', | |
| ha='center', | |
| color=color) | |
| plt.legend() | |
| plt.tight_layout() | |
| # Save plot | |
| buf = io.BytesIO() | |
| plt.savefig(buf, format='png', bbox_inches='tight', dpi=300) | |
| buf.seek(0) | |
| plot_image = Image.open(buf) | |
| plt.close() | |
| # Calculate final probabilities | |
| with torch.no_grad(): | |
| output = model(X_tensor) | |
| probs = torch.softmax(output, dim=1) | |
| pred_class = 1 if probs[0][1] > probs[0][0] else 0 | |
| pred_label = 'human' if pred_class == 1 else 'non-human' | |
| # Generate results text | |
| results_text += f"""Sequence: {header} | |
| Prediction: {pred_label} | |
| Confidence: {float(max(probs[0])):0.4f} | |
| Human probability: {float(probs[0][1]):0.4f} | |
| Non-human probability: {float(probs[0][0]):0.4f} | |
| Most influential k-mers (ranked by importance):""" | |
| for kmer in important_kmers: | |
| direction = "human" if kmer['importance'] > 0 else "non-human" | |
| results_text += f"\n {kmer['kmer']}: " | |
| results_text += f"pushes toward {direction} (impact={abs(kmer['importance']):.4f}), " | |
| results_text += f"occurrence={kmer['frequency']*100:.2f}% of sequence " | |
| if kmer['scaled'] > 0: | |
| results_text += f"(appears {abs(kmer['scaled']):.2f}σ more than average)" | |
| else: | |
| results_text += f"(appears {abs(kmer['scaled']):.2f}σ less than average)" | |
| except Exception as e: | |
| return f"Error processing sequences: {str(e)}", None | |
| return results_text, plot_image | |
| iface = gr.Interface( | |
| fn=predict, | |
| inputs=gr.File(label="Upload FASTA file", type="binary"), | |
| outputs=[gr.Textbox(label="Results"), gr.Image(label="SHAP Waterfall Plot")], | |
| title="Virus Host Classifier" | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch(share=True) |