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 | |
| 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) | |
| # Get feature importance and human probability | |
| importance, human_prob = model.get_feature_importance(X_tensor) | |
| kmer_importance = importance[0].cpu().numpy() | |
| # Scale importance values relative to the prediction | |
| kmer_importance = kmer_importance * human_prob | |
| # 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) | |
| # Create SHAP explanation | |
| # Set base_value to 0.5 (neutral prediction) | |
| # Values represent the push towards human (>0.5) or non-human (<0.5) | |
| 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 | |
| # Create waterfall plot | |
| plt.figure(figsize=(10, 6)) | |
| fig = shap.plots._waterfall.waterfall_legacy( | |
| explanation, | |
| show=False, | |
| max_display=11 # Show all features including "Others" | |
| ) | |
| plt.title(f"Impact on prediction (>0.5 pushes toward human, <0.5 toward non-human)") | |
| # 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) | |