Update app.py

app.py

@@ -231,7 +231,6 @@ def plot_all():
 
 # ---------- Prepare Data for Modeling ----------
 X = df.iloc[:, 1:].values
-y = df['Label'].values
 scaler = StandardScaler()
 X_scaled = scaler.fit_transform(X)
 pca = PCA(n_components=2)
@@ -269,11 +268,10 @@ class CNN1D(nn.Module):
             nn.Conv1d(32, 64, 3, padding=1), nn.ReLU(),
             nn.AdaptiveAvgPool1d(1),
             nn.Flatten(),
-            nn.Linear(64, num_classes)
+            nn.Linear(64, len(np.unique(y)))
         )
     def forward(self, x): return self.net(x)
-        
-num_classes = len(np.unique(np.concatenate([y_train_raw, y_test_raw])))
+
 model = CNN1D()
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.Adam(model.parameters(), lr=0.001)
@@ -350,9 +348,11 @@ with gr.Blocks() as demo:
                     X_tensor = torch.tensor(X_scaled, dtype=torch.float32).unsqueeze(1)
                     with torch.no_grad():
                         preds = torch.argmax(model(X_tensor), dim=1).numpy()
+                        preds = le.inverse_transform(preds)
                 else:
                     X_pca_input = pca.transform(scaler.transform(X_input))
                     preds = rf.predict(X_pca_input) if model_name == 'Random Forest' else dt.predict(X_pca_input)
+                    preds = le.inverse_transform(preds)
                 test_df['Predicted Label'] = preds
                 return test_df