Update app.py

app.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 import shap
 import matplotlib.pyplot as plt
 import io
-from PIL import Image  # Import PIL for image handling
+from PIL import Image
 
 class VirusClassifier(nn.Module):
     def __init__(self, input_shape: int):
@@ -30,18 +30,19 @@ class VirusClassifier(nn.Module):
         return self.network(x)
     
     def get_feature_importance(self, x):
-        """Calculate feature importance using gradient-based method"""
+        """Calculate feature importance using gradient-based method for the human class (index 1)"""
         x.requires_grad_(True)
         output = self.network(x)
-        importance = torch.zeros_like(x)
+        probs = torch.softmax(output, dim=1)
         
-        for i in range(output.shape[1]):
-            if x.grad is not None:
-                x.grad.zero_()
-            output[..., i].sum().backward(retain_graph=True)
-            importance += torch.abs(x.grad)
-            
-        return importance
+        # 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"""
@@ -111,23 +112,20 @@ def predict(file_obj):
     
     try:
         sequences = parse_fasta(text)
-        # For simplicity, process only the first sequence for plotting
         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)
         
-        with torch.no_grad():
-            output = model(X_tensor)
-            probs = torch.softmax(output, dim=1)
-        
-        importance = model.get_feature_importance(X_tensor)
+        # Get feature importance and human probability
+        importance, human_prob = model.get_feature_importance(X_tensor)
         kmer_importance = importance[0].cpu().numpy()
         
-        if np.max(np.abs(kmer_importance)) != 0:
-            kmer_importance = kmer_importance / np.max(np.abs(kmer_importance)) * 0.002
+        # 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 = [
@@ -140,9 +138,11 @@ def predict(file_obj):
             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])
@@ -150,10 +150,12 @@ def predict(file_obj):
         top_features.append("Others")
         top_values.append(others_sum)
         
-        # Set base_values and expected_value to 0 for the binary classification starting point
+        # 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.0,  
+            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]) 
@@ -161,18 +163,33 @@ def predict(file_obj):
             ]),
             feature_names=top_features
         )
-        explanation.expected_value = 0.0
-        
-        fig = shap.plots._waterfall.waterfall_legacy(explanation, show=False)
+        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()
-        fig.savefig(buf, format='png')
+        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}
@@ -181,8 +198,9 @@ 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"impact={kmer['importance']:.4f}, "
+            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)"
@@ -203,4 +221,3 @@ iface = gr.Interface(
 
 if __name__ == "__main__":
     iface.launch(share=True)
-