Update app.py

app.py

@@ -7,6 +7,7 @@ import torch.nn as nn
 import shap
 import matplotlib.pyplot as plt
 import io
+import json
 from PIL import Image
 
 class VirusClassifier(nn.Module):
@@ -183,6 +184,7 @@ def predict(file_obj):
             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
@@ -196,36 +198,67 @@ def predict(file_obj):
         explanation.expected_value = 0.5  # Start from neutral prediction
 
         # Calculate step-by-step probabilities
-        step_probs = []
         current_prob = 0.5  # Start at neutral
-        step_probs.append({"step": "Start", "probability": current_prob, "kmer": "Initial", "change": 0})
+        steps = [('Start', current_prob, 0)]
         
         # Process each k-mer contribution
-        for i, kmer in enumerate(important_kmers, 1):
+        for kmer in important_kmers:
             change = kmer['importance']
             current_prob += change
-            step_probs.append({
-                "step": str(i),
-                "probability": current_prob,
-                "kmer": kmer['kmer'],
-                "change": change
-            })
+            steps.append((kmer['kmer'], current_prob, change))
         
         # Add final "Others" contribution
         current_prob += others_sum
-        step_probs.append({
-            "step": "Others",
-            "probability": current_prob,
-            "kmer": "Others",
-            "change": 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)
         
-        # Convert to JSON for React
-        steps_json = json.dumps(step_probs, indent=2)
-        print(f"Steps data: {steps_json}")  # For debugging
+        plt.legend()
+        plt.tight_layout()
         
-        # Create visualization component
-        plot_image = None  # We'll use the React component instead
+        # 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():