Update app.py

app.py

@@ -93,7 +93,11 @@ def create_importance_plot(shap_values, kmers, top_k=10):
     """
     Create horizontal bar plot of feature importance.
     """
-    plt.style.use('seaborn-v0_8-whitegrid')
+    # Set style directly instead of using seaborn
+    plt.rcParams['figure.facecolor'] = '#ffffff'
+    plt.rcParams['axes.facecolor'] = '#ffffff'
+    plt.rcParams['axes.grid'] = True
+    plt.rcParams['grid.alpha'] = 0.3
     fig = plt.figure(figsize=(10, 8))
     
     # Sort by absolute importance
@@ -115,8 +119,13 @@ def create_contribution_plot(important_kmers, final_prob):
     """
     Create waterfall plot showing cumulative feature contributions.
     """
-    plt.style.use('seaborn-v0_8-whitegrid')
-    fig = plt.figure(figsize=(12, 6))
+    # Set style parameters
+    plt.rcParams['figure.facecolor'] = '#ffffff'
+    plt.rcParams['axes.facecolor'] = '#ffffff'
+    plt.rcParams['axes.grid'] = True
+    plt.rcParams['grid.alpha'] = 0.3
+    
+    fig, ax = plt.subplots(figsize=(12, 6))
     
     base_prob = 0.5
     cumulative = [base_prob]
@@ -126,15 +135,36 @@ def create_contribution_plot(important_kmers, final_prob):
         cumulative.append(cumulative[-1] + kmer_info['impact'])
         labels.append(kmer_info['kmer'])
     
-    plt.plot(range(len(cumulative)), cumulative, 'b-o', linewidth=2)
-    plt.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
+    # Plot cumulative line with markers
+    line = ax.plot(range(len(cumulative)), cumulative, '-o', 
+                  color='#3498db', linewidth=2, 
+                  marker='o', markersize=8, 
+                  markerfacecolor='white',
+                  markeredgecolor='#3498db',
+                  markeredgewidth=2)
+    
+    # Add reference line at 0.5
+    ax.axhline(y=0.5, color='#95a5a6', linestyle='--', alpha=0.5)
+    
+    # Customize plot
+    ax.set_xticks(range(len(labels)))
+    ax.set_xticklabels(labels, rotation=45, ha='right')
+    ax.set_ylim(0, 1)
+    ax.grid(True, axis='y', linestyle='--', alpha=0.3)
+    ax.set_title('Cumulative Feature Contributions')
+    ax.set_ylabel('Probability of Human Origin')
     
-    plt.xticks(range(len(labels)), labels, rotation=45)
-    plt.ylim(0, 1)
-    plt.grid(True, alpha=0.3)
-    plt.title('Cumulative Feature Contributions')
-    plt.ylabel('Probability of Human Origin')
+    # Add value labels
+    for i, prob in enumerate(cumulative):
+        ax.annotate(f'{prob:.3f}',
+                   (i, prob),
+                   xytext=(0, 10),
+                   textcoords='offset points',
+                   ha='center',
+                   va='bottom')
     
+    # Adjust layout to prevent label cutoff
+    plt.tight_layout()
     return fig
 
 def predict(file_obj, top_kmers=10, fasta_text=""):
@@ -165,7 +195,8 @@ def predict(file_obj, top_kmers=10, fasta_text=""):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     try:
         model = VirusClassifier(256).to(device)
-        model.load_state_dict(torch.load('model.pt', map_location=device))
+        # Load model weights safely
+        model.load_state_dict(torch.load('model.pt', map_location=device, weights_only=True))
         scaler = joblib.load('scaler.pkl')
     except Exception as e:
         return f"Error loading model: {str(e)}", None, None