Update app.py

app.py

@@ -12,7 +12,7 @@ from sklearn.manifold import TSNE
 import umap.umap_ as umap
 import shap
 import joblib
-from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_curve, auc, mean_squared_error, mean_absolute_error, r2_score, classification_report
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_curve, auc, mean_squared_error, mean_absolute_error, r2_score, classification_report, silhouette_score
 from sklearn.pipeline import Pipeline
 from sklearn.compose import ColumnTransformer
 from sklearn.impute import SimpleImputer
@@ -20,15 +20,66 @@ from sklearn.cluster import KMeans, DBSCAN
 from sklearn.mixture import GaussianMixture
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
 from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.svm import SVC, SVR
 from xgboost import XGBClassifier, XGBRegressor
 import matplotlib.pyplot as plt
 from io import BytesIO
 import time
+from PIL import Image
+import zipfile
+import os
 
 # Set page config
 st.set_page_config(page_title="Neural-Vision Enhanced", layout="wide")
 
-# Helper Functions
+# Helper Functions for Image Processing
+def preprocess_image(image_path, target_size=(224, 224)):
+    """Preprocess an image by resizing and normalizing it."""
+    img = Image.open(image_path).convert("RGB")
+    img = img.resize(target_size)
+    img_array = np.array(img) / 255.0  # Normalize pixel values to [0, 1]
+    return img_array
+
+def load_image_dataset(zip_path, target_size=(224, 224), problem_type="Classification"):
+    """Load and preprocess an image dataset from a zip file."""
+    # Extract zip file to a temporary directory
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall('temp_images')
+    
+    if problem_type == "Classification":
+        image_paths = []
+        labels = []
+        class_names = sorted(os.listdir('temp_images'))
+        for label, class_name in enumerate(class_names):
+            class_dir = os.path.join('temp_images', class_name)
+            if os.path.isdir(class_dir):
+                for img_name in os.listdir(class_dir):
+                    image_path = os.path.join(class_dir, img_name)
+                    if os.path.isfile(image_path):
+                        image_paths.append(image_path)
+                        labels.append(label)
+        images = [preprocess_image(path, target_size) for path in image_paths]
+        images = np.array(images)
+        labels = np.array(labels)
+        data = (images, labels, class_names)
+    else:  # Compression or Clustering
+        image_dir = 'temp_images'
+        image_paths = [os.path.join(image_dir, img_name) for img_name in os.listdir(image_dir) if os.path.isfile(os.path.join(image_dir, img_name))]
+        images = [preprocess_image(path, target_size) for path in image_paths]
+        images = np.array(images)
+        data = (images, None, None)
+    
+    # Clean up temporary directory
+    for root, dirs, files in os.walk('temp_images', topdown=False):
+        for name in files:
+            os.remove(os.path.join(root, name))
+        for name in dirs:
+            os.rmdir(os.path.join(root, name))
+    os.rmdir('temp_images')
+    
+    return data
+
+# Model Building Functions
 def get_model_config(model_type, problem_type):
     configs = {
         "Random Forest": {
@@ -54,6 +105,11 @@ def get_model_config(model_type, problem_type):
         "Linear Regression": {
             "Regression": {"model_class": LinearRegression, "params": {}, "grid_params": {}}
         },
+        "SVM": {
+            "Regression": {"model_class": SVR, "params": {"kernel": "rbf"}, "grid_params": {"C": [0.1, 1, 10], "gamma": ["scale", "auto"]}},
+            "Binary Classification": {"model_class": SVC, "params": {"kernel": "rbf", "random_state": 42}, "grid_params": {"C": [0.1, 1, 10], "gamma": ["scale", "auto"]}},
+            "Multi-Class": {"model_class": SVC, "params": {"kernel": "rbf", "random_state": 42}, "grid_params": {"C": [0.1, 1, 10], "gamma": ["scale", "auto"]}}
+        },
         "K-Means": {
             "Clustering": {"model_class": KMeans, "params": {"n_clusters": 3, "random_state": 42},
                            "grid_params": {"n_clusters": [2, 3, 4, 5]}}
@@ -68,85 +124,130 @@ def get_model_config(model_type, problem_type):
         }
     }
     return configs.get(model_type, {}).get(problem_type, {"model_class": None, "params": {}, "grid_params": {}})
-    
+
 def preprocess_data(X_train, X_test, numerical_features, categorical_features):
-    # Define the numeric and categorical transformers
     numeric_transformer = Pipeline(steps=[
         ('imputer', SimpleImputer(strategy='mean')),
         ('scaler', StandardScaler())])
-
     categorical_transformer = Pipeline(steps=[
         ('imputer', SimpleImputer(strategy='most_frequent')),
         ('onehot', OneHotEncoder(handle_unknown='ignore', sparse_output=False))])
-
-    # Combine transformers into a ColumnTransformer
     preprocessor = ColumnTransformer(
         transformers=[
             ('num', numeric_transformer, numerical_features),
             ('cat', categorical_transformer, categorical_features)],
-        remainder='passthrough') # Handle unseen columns
-
-    # Fit and transform the training data
+        remainder='drop')
     X_train_processed = preprocessor.fit_transform(X_train)
-
-    # Transform the test data
     X_test_processed = preprocessor.transform(X_test)
-
-    # Get feature names after one-hot encoding
     if categorical_features:
-        # Access the fitted OneHotEncoder
         onehot_encoder = preprocessor.named_transformers_['cat'].named_steps['onehot']
         categorical_feature_names = onehot_encoder.get_feature_names_out(categorical_features)
         feature_names = numerical_features + list(categorical_feature_names)
     else:
         feature_names = numerical_features
-
     return X_train_processed, X_test_processed, feature_names, preprocessor
-def build_neural_network(input_shape, output_units, problem_type, layers_config, optimizer_name="Adam", loss_function="mse", metrics=["accuracy"]):
-    try:
-        model = keras.Sequential()
-        model.add(keras.layers.InputLayer(input_shape=input_shape))
+
+def build_neural_network(input_shape, output_units, problem_type, layers_config, optimizer_name="Adam", learning_rate=0.001):
+    model = keras.Sequential()
+    model.add(keras.layers.InputLayer(input_shape=input_shape))
+    for layer in layers_config:
+        if layer['type'] == 'dense':
+            model.add(keras.layers.Dense(layer['units'], activation=layer['activation']))
+        elif layer['type'] == 'dropout':
+            model.add(keras.layers.Dropout(layer['rate']))
+        elif layer['type'] == 'conv2d':
+            model.add(keras.layers.Conv2D(layer['filters'], tuple(layer['kernel_size']), activation=layer['activation'], padding='same'))
+        elif layer['type'] == 'maxpooling2d':
+            model.add(keras.layers.MaxPooling2D(pool_size=tuple(layer['pool_size'])))
+        elif layer['type'] == 'flatten':
+            model.add(keras.layers.Flatten())
+    if problem_type == "Regression":
+        model.add(keras.layers.Dense(1))
+        loss_function = "mse"
+        metrics = ["mse"]
+    elif problem_type == "Binary Classification":
+        model.add(keras.layers.Dense(1, activation='sigmoid'))
+        loss_function = "binary_crossentropy"
+        metrics = ["accuracy"]
+    elif problem_type == "Multi-Class" or problem_type == "Image Classification":
+        model.add(keras.layers.Dense(output_units, activation='softmax'))
+        loss_function = "sparse_categorical_crossentropy" if problem_type == "Image Classification" else "categorical_crossentropy"
+        metrics = ["accuracy"]
+    else:
+        raise ValueError("Unsupported problem type")
+    optimizer = {"Adam": keras.optimizers.Adam, "SGD": keras.optimizers.SGD, "RMSprop": keras.optimizers.RMSprop}.get(optimizer_name)(learning_rate=learning_rate)
+    model.compile(optimizer=optimizer, loss=loss_function, metrics=metrics)
+    return model
+
+def build_autoencoder(input_shape, encoding_dim, layers_config, autoencoder_type="Standard", optimizer_name="Adam", learning_rate=0.001):
+    if autoencoder_type == "Variational":
+        inputs = keras.layers.Input(shape=input_shape)
+        x = keras.layers.Flatten()(inputs) if len(input_shape) > 1 else inputs
         for layer in layers_config:
             if layer['type'] == 'dense':
-                model.add(keras.layers.Dense(layer['units'], activation=layer['activation']))
+                x = keras.layers.Dense(layer['units'], activation=layer['activation'])(x)
+        z_mean = keras.layers.Dense(encoding_dim, name='z_mean')(x)
+        z_log_var = keras.layers.Dense(encoding_dim, name='z_log_var')(x)
+        
+        def sampling(args):
+            z_mean, z_log_var = args
+            epsilon = keras.backend.random_normal(shape=(keras.backend.shape(z_mean)[0], encoding_dim))
+            return z_mean + keras.backend.exp(0.5 * z_log_var) * epsilon
+        
+        z = keras.layers.Lambda(sampling, name='z')([z_mean, z_log_var])
+        encoder = keras.Model(inputs, [z_mean, z_log_var, z], name='encoder')
+        
+        decoder_input = keras.layers.Input(shape=(encoding_dim,))
+        x = decoder_input
+        for layer in reversed(layers_config):
+            if layer['type'] == 'dense':
+                x = keras.layers.Dense(layer['units'], activation=layer['activation'])(x)
+        x = keras.layers.Dense(np.prod(input_shape), activation='sigmoid')(x)
+        outputs = keras.layers.Reshape(input_shape)(x) if len(input_shape) > 1 else x
+        decoder = keras.Model(decoder_input, outputs, name='decoder')
+        
+        vae_outputs = decoder(encoder(inputs)[2])
+        autoencoder = keras.Model(inputs, vae_outputs, name='vae')
+        
+        reconstruction_loss = keras.losses.binary_crossentropy(keras.backend.flatten(inputs), keras.backend.flatten(vae_outputs))
+        reconstruction_loss *= np.prod(input_shape)
+        kl_loss = 1 + z_log_var - keras.backend.square(z_mean) - keras.backend.exp(z_log_var)
+        kl_loss = keras.backend.sum(kl_loss, axis=-1) * -0.5
+        vae_loss = keras.backend.mean(reconstruction_loss + kl_loss)
+        autoencoder.add_loss(vae_loss)
+    else:  # Standard or Denoising
+        encoder = keras.Sequential([keras.layers.InputLayer(input_shape=input_shape)])
+        for layer in layers_config:
+            if layer['type'] == 'dense':
+                encoder.add(keras.layers.Dense(layer['units'], activation=layer['activation']))
             elif layer['type'] == 'dropout':
-                model.add(keras.layers.Dropout(layer['rate']))
-            elif layer['type'] == 'conv2d':
-                model.add(keras.layers.Conv2D(layer['filters'], tuple(layer['kernel_size']), activation=layer['activation']))
-            elif layer['type'] == 'lstm':
-                model.add(keras.layers.LSTM(layer['units'], activation=layer['activation']))
-            elif layer['type'] == 'maxpooling2d':
-                model.add(keras.layers.MaxPooling2D(pool_size=tuple(layer['pool_size'])))
-            elif layer['type'] == 'flatten':
-                model.add(keras.layers.Flatten())
-
-        if problem_type == "Regression":
-            model.add(keras.layers.Dense(1))
-        elif problem_type == "Binary Classification":
-            model.add(keras.layers.Dense(1, activation='sigmoid'))
-        elif problem_type == "Multi-Class":
-            model.add(keras.layers.Dense(output_units, activation='softmax'))
-
-        if optimizer_name == "Adam":
-            optimizer = keras.optimizers.Adam()
-        elif optimizer_name == "SGD":
-            optimizer = keras.optimizers.SGD()
-        elif optimizer_name == "RMSprop":
-            optimizer = keras.optimizers.RMSprop()
-
-        model.compile(optimizer=optimizer, loss=loss_function, metrics=metrics)
-        return model
-    except Exception as e:
-        st.error("Failed to build neural network. Check Debug Log for details.")
-        #log_error("Error in build_neural_network", e) # ADDED
-        raise
-
+                encoder.add(keras.layers.Dropout(layer['rate']))
+        encoder.add(keras.layers.Dense(encoding_dim, activation='relu', name='encoded'))
+        
+        decoder = keras.Sequential([keras.layers.InputLayer(input_shape=(encoding_dim,))])
+        for layer in reversed(layers_config):
+            if layer['type'] == 'dense':
+                decoder.add(keras.layers.Dense(layer['units'], activation=layer['activation']))
+        decoder.add(keras.layers.Dense(np.prod(input_shape), activation='sigmoid'))
+        decoder.add(keras.layers.Reshape(input_shape) if len(input_shape) > 1 else keras.layers.Lambda(lambda x: x))
+        
+        autoencoder_input = keras.layers.Input(shape=input_shape)
+        encoded = encoder(autoencoder_input)
+        decoded = decoder(encoded)
+        autoencoder = keras.Model(autoencoder_input, decoded)
+    
+    optimizer = {"Adam": keras.optimizers.Adam, "SGD": keras.optimizers.SGD, "RMSprop": keras.optimizers.RMSprop}.get(optimizer_name)(learning_rate=learning_rate)
+    autoencoder.compile(optimizer=optimizer, loss='mse', metrics=['mse'])
+    return autoencoder, encoder, decoder
 
-def train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size, problem_type, do_grid_search=False, params=None, grid_params=None, training_placeholder=None):
+def train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size, problem_type, input_data=None, target_data=None, do_grid_search=False, params=None, grid_params=None, training_placeholder=None):
     start_time = time.time()
     history = None
+    if input_data is None:
+        input_data = X_train
+    if target_data is None:
+        target_data = y_train if y_train is not None else X_train
     if isinstance(model, keras.Model):
-        # Define a callback to update Streamlit during training
         class StreamlitCallback(keras.callbacks.Callback):
             def __init__(self, placeholder):
                 super().__init__()
@@ -157,31 +258,28 @@ def train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size, pro
                 self.epoch_data.append(logs)
                 df = pd.DataFrame(self.epoch_data)
                 fig = px.line(df, x=df.index, y=['loss', 'val_loss'], labels={'index': 'Epoch', 'value': 'Loss'})
-                fig.add_trace(go.Scatter(x=df.index, y=df['accuracy'], mode='lines', name='accuracy'))
-                fig.add_trace(go.Scatter(x=df.index, y=df['val_accuracy'], mode='lines', name='val_accuracy'))
-
+                metric_name = 'mse' if 'mse' in logs else 'accuracy'
+                fig.add_trace(go.Scatter(x=df.index, y=df[metric_name], mode='lines', name=metric_name))
+                fig.add_trace(go.Scatter(x=df.index, y=df[f'val_{metric_name}'], mode='lines', name=f'val_{metric_name}'))
                 self.placeholder.plotly_chart(fig)
 
         streamlit_callback = StreamlitCallback(training_placeholder)
-        history = model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size,
-                              validation_data=(X_test, y_test), verbose=0,
-                              callbacks=[streamlit_callback])
-
+        history = model.fit(input_data, target_data, epochs=epochs, batch_size=batch_size,
+                            validation_data=(X_test, y_test if y_test is not None else X_test), verbose=0,
+                            callbacks=[streamlit_callback])
     else:
         if do_grid_search and grid_params:
-            grid_search = GridSearchCV(model, grid_params, cv=3, n_jobs=-1, scoring='accuracy' if problem_type in ["Binary Classification", "Multi-Class"] else 'neg_mean_squared_error') # Added scoring
+            grid_search = GridSearchCV(model, grid_params, cv=3, n_jobs=-1, scoring='accuracy' if problem_type in ["Binary Classification", "Multi-Class"] else 'neg_mean_squared_error')
             grid_search.fit(X_train, y_train)
             model = grid_search.best_estimator_
-            st.write("Best parameters found by Grid Search:", grid_search.best_params_) # Print best params
+            st.write("Best parameters found by Grid Search:", grid_search.best_params_)
         else:
             model.set_params(**params)
             model.fit(X_train, y_train)
-            history = None # Reset to none here
-
     training_time = time.time() - start_time
     return history, model, training_time
 
-def evaluate_model(model, X_test, y_test, problem_type, le=None):
+def evaluate_model(model, X_test, y_test, problem_type, le=None, encoder=None):
     y_pred = model.predict(X_test)
     metrics = {}
     if problem_type == "Regression":
@@ -190,9 +288,13 @@ def evaluate_model(model, X_test, y_test, problem_type, le=None):
         metrics['rmse'] = np.sqrt(metrics['mse'])
         metrics['r2'] = r2_score(y_test, y_pred)
         return metrics, y_pred.flatten()
-    elif problem_type in ["Binary Classification", "Multi-Class"]:
-        y_pred_classes = (y_pred > 0.5).astype(int).flatten() if problem_type == "Binary Classification" else np.argmax(y_pred, axis=1)
-        y_test_classes = y_test if problem_type == "Binary Classification" else np.argmax(y_test, axis=1)
+    elif problem_type in ["Binary Classification", "Multi-Class", "Image Classification"]:
+        if problem_type == "Image Classification":
+            y_pred_classes = np.argmax(y_pred, axis=1)
+            y_test_classes = y_test
+        else:
+            y_pred_classes = (y_pred > 0.5).astype(int).flatten() if problem_type == "Binary Classification" else np.argmax(y_pred, axis=1)
+            y_test_classes = y_test if problem_type == "Binary Classification" else np.argmax(y_test, axis=1)
         metrics['accuracy'] = accuracy_score(y_test_classes, y_pred_classes)
         metrics['precision'] = precision_score(y_test_classes, y_pred_classes, average='weighted', zero_division=0)
         metrics['recall'] = recall_score(y_test_classes, y_pred_classes, average='weighted', zero_division=0)
@@ -200,7 +302,16 @@ def evaluate_model(model, X_test, y_test, problem_type, le=None):
         return metrics, y_pred_classes
     elif problem_type == "Clustering":
         labels = model.labels_ if hasattr(model, 'labels_') else model.predict(X_test)
-        return {"n_clusters": len(np.unique(labels))}, labels
+        metrics["n_clusters"] = len(np.unique(labels))
+        if len(np.unique(labels)) > 1:
+            metrics["silhouette"] = silhouette_score(X_test, labels)
+        return metrics, labels
+    elif problem_type == "Compression":
+        metrics['mse'] = mean_squared_error(X_test, y_pred)
+        metrics['mae'] = mean_absolute_error(X_test, y_pred)
+        metrics['rmse'] = np.sqrt(metrics['mse'])
+        compressed_data = encoder.predict(X_test) if encoder else None
+        return metrics, y_pred, compressed_data
 
 def save_model(model, preprocessor, features, target, problem_type, filename="model.pkl"):
     model_data = {
@@ -229,42 +340,63 @@ with st.sidebar:
     st.markdown("Your AI-powered model toolbox.")
     st.markdown("---")
     app_mode = st.selectbox("Navigation", ["Data Upload", "Model Training", "Validation & Exploration"])
+    data_type = st.selectbox("Data Type", ["Tabular", "Image"])
     st.markdown("---")
-    st.markdown("**Dependencies**: `tensorflow`, `shap`, `umap-learn`, `joblib`, `scikit-learn`, `plotly`, `xgboost`")
-    st.markdown("Created by Calvin Allen-Crawford | v1.2 | © 2025")
+    st.markdown("**Dependencies**: `tensorflow`, `shap`, `umap-learn`, `joblib`, `scikit-learn`, `plotly`, `xgboost`, `pillow`")
+    st.markdown("Created by Calvin Allen-Crawford | v1.3 | © 2025")
 
 # Main App Sections
 if app_mode == "Data Upload":
     st.title("📤 Data Upload")
     col1, col2, col3 = st.columns([1, 2, 1])
     with col2:
-        uploaded_file = st.file_uploader("Upload CSV Dataset", type=["csv"])
-
-    if uploaded_file:
-        df = pd.read_csv(uploaded_file)
-        st.session_state.df = df
-        st.write("---")
-        st.subheader("Dataset Preview")
-        st.dataframe(df.head(10))
-        st.write("---")
-        st.subheader("Statistics")
-        col1, col2, col3 = st.columns(3)
-        with col1: st.metric("Rows", df.shape[0])
-        with col2: st.metric("Columns", df.shape[1])
-        with col3: st.metric("Missing Values", df.isna().sum().sum())
+        if data_type == "Tabular":
+            uploaded_file = st.file_uploader("Upload CSV Dataset", type=["csv"])
+            if uploaded_file:
+                df = pd.read_csv(uploaded_file)
+                st.session_state.df = df
+                st.write("---")
+                st.subheader("Dataset Preview")
+                st.dataframe(df.head(10))
+                st.write("---")
+                st.subheader("Statistics")
+                col1, col2, col3 = st.columns(3)
+                with col1: st.metric("Rows", df.shape[0])
+                with col2: st.metric("Columns", df.shape[1])
+                with col3: st.metric("Missing Values", df.isna().sum().sum())
+        else:  # Image
+            uploaded_file = st.file_uploader("Upload Zip File with Images", type=["zip"])
+            if uploaded_file:
+                problem_type = st.selectbox("Problem Type for Image Data", ["Image Classification", "Compression", "Clustering"])
+                images, labels, class_names = load_image_dataset(uploaded_file, problem_type=problem_type)
+                st.session_state.images = images
+                st.session_state.labels = labels
+                st.session_state.class_names = class_names if problem_type == "Image Classification" else None
+                st.write(f"Loaded {len(images)} images.")
+                if problem_type == "Image Classification":
+                    st.write(f"Classes: {class_names}")
+                st.image(images[:5], caption=["Sample " + str(i+1) for i in range(min(5, len(images)))], width=100)
 
 elif app_mode == "Model Training":
     st.title("🧠 Model Training")
-    if 'df' not in st.session_state:
-        st.warning("Please upload a dataset first.")
+    if data_type == "Tabular" and 'df' not in st.session_state:
+        st.warning("Please upload a tabular dataset first.")
+        st.stop()
+    elif data_type == "Image" and 'images' not in st.session_state:
+        st.warning("Please upload an image dataset first.")
         st.stop()
 
-    df = st.session_state.df
-    problem_type = st.selectbox("Problem Type", ["Regression", "Binary Classification", "Multi-Class", "Clustering"])
-    features = st.multiselect("Select Features", df.columns)
-    target = st.selectbox("Select Target", df.columns) if problem_type != "Clustering" else None
+    if data_type == "Tabular":
+        df = st.session_state.df
+        problem_type = st.selectbox("Problem Type", ["Regression", "Binary Classification", "Multi-Class", "Clustering", "Compression"])
+        features = st.multiselect("Select Features", df.columns)
+        target = st.selectbox("Select Target", df.columns) if problem_type not in ["Clustering", "Compression"] else None
+    else:
+        problem_type = st.selectbox("Problem Type", ["Image Classification", "Compression", "Clustering"])
+        features = ["images"]
+        target = "labels" if problem_type == "Image Classification" else None
 
-    if problem_type != "Clustering" and target:
+    if problem_type not in ["Clustering", "Compression"] and data_type == "Tabular" and target:
         unique_target_values = df[target].nunique()
         if problem_type == "Binary Classification" and unique_target_values != 2:
             st.error("Binary Classification requires exactly 2 unique target values.")
@@ -277,17 +409,20 @@ elif app_mode == "Model Training":
             st.stop()
 
     model_types = {
-        "Regression": ["Neural Network", "Random Forest", "XGBoost", "Linear Regression"],
-        "Binary Classification": ["Neural Network", "Random Forest", "XGBoost", "Logistic Regression"],
-        "Multi-Class": ["Neural Network", "Random Forest", "XGBoost"],
-        "Clustering": ["K-Means", "DBSCAN", "Gaussian Mixture"]
+        "Regression": ["Neural Network", "Random Forest", "XGBoost", "Linear Regression", "SVM"],
+        "Binary Classification": ["Neural Network", "Random Forest", "XGBoost", "Logistic Regression", "SVM"],
+        "Multi-Class": ["Neural Network", "Random Forest", "XGBoost", "SVM"],
+        "Clustering": ["K-Means", "DBSCAN", "Gaussian Mixture"],
+        "Compression": ["Autoencoder"],
+        "Image Classification": ["Neural Network", "SVM"]
     }[problem_type]
     model_type = st.selectbox("Model Type", model_types)
 
     if model_type == "Neural Network":
         st.subheader("Neural Network Configuration")
+        optimizer_name = st.selectbox("Optimizer", ["Adam", "SGD", "RMSprop"])
         layers_config = st.session_state.get('layers_config', [])
-        layer_type = st.selectbox("Layer Type", ["Dense", "Dropout"])
+        layer_type = st.selectbox("Layer Type", ["Dense", "Dropout"] if data_type == "Tabular" else ["Conv2D", "MaxPooling2D", "Flatten", "Dense", "Dropout"])
         if layer_type == "Dense":
             units = st.number_input("Units", min_value=1, value=64)
             activation = st.selectbox("Activation", ["relu", "sigmoid", "tanh"])
@@ -297,14 +432,50 @@ elif app_mode == "Model Training":
             rate = st.number_input("Dropout Rate", 0.0, 1.0, 0.2)
             if st.button("Add Layer"):
                 layers_config.append({"type": "dropout", "rate": rate})
-
+        elif layer_type == "Conv2D":
+            filters = st.number_input("Filters", min_value=1, value=32)
+            kernel_size = st.multiselect("Kernel Size", options=[1, 3, 5], default=[3])
+            activation = st.selectbox("Activation", ["relu", "sigmoid", "tanh"])
+            if st.button("Add Layer"):
+                layers_config.append({"type": "conv2d", "filters": filters, "kernel_size": kernel_size, "activation": activation})
+        elif layer_type == "MaxPooling2D":
+            pool_size = st.multiselect("Pool Size", options=[2, 3], default=[2])
+            if st.button("Add Layer"):
+                layers_config.append({"type": "maxpooling2d", "pool_size": pool_size})
+        elif layer_type == "Flatten":
+            if st.button("Add Layer"):
+                layers_config.append({"type": "flatten"})
         if layers_config:
             st.write("Current Layers:", layers_config)
             if st.button("Clear Layers"):
                 layers_config.clear()
                 st.session_state.layers_config = layers_config
                 st.rerun()
-
+        st.session_state.layers_config = layers_config
+    elif model_type == "Autoencoder":
+        st.subheader("Autoencoder Configuration")
+        encoding_dim = st.number_input("Encoding Dimension", min_value=1, value=32)
+        optimizer_name = st.selectbox("Optimizer", ["Adam", "SGD", "RMSprop"])
+        autoencoder_type = st.selectbox("Autoencoder Type", ["Standard", "Variational", "Denoising"])
+        if autoencoder_type == "Denoising":
+            noise_level = st.slider("Noise Level", 0.0, 1.0, 0.1)
+        layers_config = st.session_state.get('layers_config', [])
+        layer_type = st.selectbox("Layer Type (Encoder)", ["Dense", "Dropout"])
+        if layer_type == "Dense":
+            units = st.number_input("Units", min_value=1, value=64)
+            activation = st.selectbox("Activation", ["relu", "sigmoid", "tanh"])
+            if st.button("Add Layer"):
+                layers_config.append({"type": "dense", "units": units, "activation": activation})
+        elif layer_type == "Dropout":
+            rate = st.number_input("Dropout Rate", 0.0, 1.0, 0.2)
+            if st.button("Add Layer"):
+                layers_config.append({"type": "dropout", "rate": rate})
+        if layers_config:
+            st.write("Current Encoder Layers:", layers_config)
+            if st.button("Clear Layers"):
+                layers_config.clear()
+                st.session_state.layers_config = layers_config
+                st.rerun()
         st.session_state.layers_config = layers_config
     else:
         st.subheader("Model Hyperparameters")
@@ -312,15 +483,8 @@ elif app_mode == "Model Training":
         params = {}
         for param_name, param_values in config["grid_params"].items():
             if isinstance(param_values[0], (int, float)) and len(param_values) > 2:
-                slider_value = st.slider(
-                    param_name,
-                    min_value=float(min(param_values)),
-                    max_value=float(max(param_values)),
-                    value=float(param_values[1])
-                )
-                # CAST TO INT FOR INTEGER PARAMETERS
-                if param_name in {'n_estimators', 'n_clusters', 'min_samples',
-                                'n_components', 'max_depth'}:
+                slider_value = st.slider(param_name, min_value=float(min(param_values)), max_value=float(max(param_values)), value=float(param_values[1]))
+                if param_name in {'n_estimators', 'n_clusters', 'min_samples', 'n_components', 'max_depth'}:
                     params[param_name] = int(slider_value)
                 else:
                     params[param_name] = slider_value
@@ -329,52 +493,75 @@ elif app_mode == "Model Training":
         do_grid_search = st.checkbox("Use Grid Search for Tuning", value=False)
 
     col1, col2, col3 = st.columns(3)
-    with col1: epochs = st.number_input("Epochs", min_value=1, value=10) if model_type == "Neural Network" else 10
-    with col2: batch_size = st.number_input("Batch Size", min_value=1, value=32) if model_type == "Neural Network" else 32
-    with col3: learning_rate = st.number_input("Learning Rate", min_value=0.0, value=0.001, step=0.0001) if model_type == "Neural Network" else 0.001
+    with col1: epochs = st.number_input("Epochs", min_value=1, value=10) if model_type in ["Neural Network", "Autoencoder"] else 10
+    with col2: batch_size = st.number_input("Batch Size", min_value=1, value=32) if model_type in ["Neural Network", "Autoencoder"] else 32
+    with col3: learning_rate = st.number_input("Learning Rate", min_value=0.0, value=0.001, step=0.0001) if model_type in ["Neural Network", "Autoencoder"] else 0.001
 
-    uploaded_model = st.file_uploader("Upload Pre-trained Model (.h5)", type=["h5"]) if model_type == "Neural Network" else None
+    uploaded_model = st.file_uploader("Upload Pre-trained Model (.h5)", type=["h5"]) if model_type in ["Neural Network", "Autoencoder"] else None
     base_model = keras.models.load_model(uploaded_model) if uploaded_model else None
 
     if st.button("Train Model"):
         with st.spinner("Preparing data..."):
-            X = df[features]
-            y = df[target] if problem_type != "Clustering" else None
-            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) if problem_type != "Clustering" else (X, X.copy(), None, None)
-            numerical_features = X.select_dtypes(include=np.number).columns.tolist()
-            categorical_features = X.select_dtypes(exclude=np.number).columns.tolist()
-            X_train_processed, X_test_processed, feature_names, preprocessor = preprocess_data(X_train, X_test, numerical_features, categorical_features)
-
-            le = None
-            if problem_type in ["Binary Classification", "Multi-Class"] and y is not None:
-                le = LabelEncoder()
-                y_train = le.fit_transform(y_train)
-                y_test = le.transform(y_test)
-                if problem_type == "Multi-Class":
-                    y_train = tf.keras.utils.to_categorical(y_train)
-                    y_test = tf.keras.utils.to_categorical(y_test)
+            if data_type == "Tabular":
+                X = df[features]
+                y = df[target] if problem_type not in ["Clustering", "Compression"] else None
+                X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) if problem_type not in ["Clustering", "Compression"] else (X, X.copy(), None, None)
+                numerical_features = X.select_dtypes(include=np.number).columns.tolist()
+                categorical_features = X.select_dtypes(exclude=np.number).columns.tolist()
+                X_train_processed, X_test_processed, feature_names, preprocessor = preprocess_data(X_train, X_test, numerical_features, categorical_features)
+                le = None
+                if problem_type in ["Binary Classification", "Multi-Class"] and y is not None:
+                    le = LabelEncoder()
+                    y_train = le.fit_transform(y_train)
+                    y_test = le.transform(y_test)
+                    if problem_type == "Multi-Class":
+                        y_train = tf.keras.utils.to_categorical(y_train)
+                        y_test = tf.keras.utils.to_categorical(y_test)
+            else:  # Image
+                X_train, X_test, y_train, y_test = train_test_split(st.session_state.images, st.session_state.labels, test_size=0.2, random_state=42) if problem_type == "Image Classification" else (st.session_state.images, st.session_state.images.copy(), None, None)
+                preprocessor = None
+                feature_names = ["image_features"]
+                le = None
 
         with st.spinner("Training model..."):
-            training_placeholder = st.empty()  # Placeholder for real-time training updates
+            training_placeholder = st.empty()
             if model_type == "Neural Network":
                 if not layers_config and not base_model:
                     st.error("Please add layers or upload a pre-trained model.")
                     st.stop()
-                model = base_model if base_model else build_neural_network(X_train_processed.shape[1:], y_train.shape[1] if problem_type == "Multi-Class" else 1,
-                                                problem_type, layers_config, "Adam", learning_rate)
-                history, model, training_time = train_model(model, X_train_processed, y_train, X_test_processed, y_test, epochs, batch_size, problem_type,
-                                                             training_placeholder=training_placeholder)
+                input_shape = (X_train_processed.shape[1],) if data_type == "Tabular" else (224, 224, 3)
+                output_units = len(st.session_state.class_names) if problem_type == "Image Classification" else (y_train.shape[1] if problem_type == "Multi-Class" else 1)
+                model = base_model if base_model else build_neural_network(input_shape, output_units, problem_type, layers_config, optimizer_name, learning_rate)
+                history, model, training_time = train_model(model, X_train_processed if data_type == "Tabular" else X_train, y_train, 
+                                                            X_test_processed if data_type == "Tabular" else X_test, y_test, 
+                                                            epochs, batch_size, problem_type, training_placeholder=training_placeholder)
+            elif model_type == "Autoencoder":
+                if not layers_config and not base_model:
+                    st.error("Please add layers to the encoder or upload a pre-trained model.")
+                    st.stop()
+                input_shape = (X_train_processed.shape[1],) if data_type == "Tabular" else (224, 224, 3)
+                model, encoder, decoder = (base_model, None, None) if base_model else build_autoencoder(input_shape, encoding_dim, layers_config, autoencoder_type, optimizer_name, learning_rate)
+                if autoencoder_type == "Denoising":
+                    X_train_noisy = X_train_processed + noise_level * np.random.normal(size=X_train_processed.shape) if data_type == "Tabular" else X_train + noise_level * np.random.normal(size=X_train.shape)
+                    input_data = X_train_noisy
+                    target_data = X_train_processed if data_type == "Tabular" else X_train
+                else:
+                    input_data = X_train_processed if data_type == "Tabular" else X_train
+                    target_data = X_train_processed if data_type == "Tabular" else X_train
+                history, model, training_time = train_model(model, X_train_processed if data_type == "Tabular" else X_train, None, 
+                                                            X_test_processed if data_type == "Tabular" else X_test, None, 
+                                                            epochs, batch_size, problem_type, input_data=input_data, target_data=target_data, 
+                                                            training_placeholder=training_placeholder)
+                st.session_state.encoder = encoder
+                st.session_state.decoder = decoder
             else:
                 config = get_model_config(model_type, problem_type)
                 model = config['model_class'](**config['params'])
-                history, model, training_time = train_model(model, X_train_processed, y_train, X_test_processed, y_test, epochs, batch_size, problem_type,
-                                                            do_grid_search, params, config['grid_params'])
-                st.subheader("Training Metrics")
-                if history:
-                  fig = px.line(x=range(len(history.history['loss'])), y=history.history['loss'], labels={'x':'Epoch', 'y':'Loss'})
-                  st.plotly_chart(fig)
-                else:
-                  st.write("No history available for this model type.")
+                X_train_flat = X_train_processed if data_type == "Tabular" else X_train.reshape(X_train.shape[0], -1)
+                X_test_flat = X_test_processed if data_type == "Tabular" else X_test.reshape(X_test.shape[0], -1)
+                history, model, training_time = train_model(model, X_train_flat, y_train, X_test_flat, y_test, epochs, batch_size, problem_type,
+                                                            do_grid_search=do_grid_search, params=params, grid_params=config['grid_params'], 
+                                                            training_placeholder=training_placeholder)
 
             st.session_state.model = model
             st.session_state.preprocessor = preprocessor
@@ -390,115 +577,153 @@ elif app_mode == "Model Training":
 
 elif app_mode == "Validation & Exploration":
     st.title("🔍 Validation & Exploration")
-    if 'model' not in st.session_state or 'df' not in st.session_state:
-        st.warning("Please upload a dataset and train a model first.")
+    if data_type == "Tabular" and ('model' not in st.session_state or 'df' not in st.session_state):
+        st.warning("Please upload a tabular dataset and train a model first.")
+        st.stop()
+    elif data_type == "Image" and ('model' not in st.session_state or 'images' not in st.session_state):
+        st.warning("Please upload an image dataset and train a model first.")
         st.stop()
 
-    df = st.session_state.df
+    if data_type == "Tabular":
+        df = st.session_state.df
+        X = df[st.session_state.features]
+        y = df[st.session_state.target] if st.session_state.problem_type not in ["Clustering", "Compression"] else None
+        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) if st.session_state.problem_type not in ["Clustering", "Compression"] else (X, X.copy(), None, None)
+        numerical_features = X.select_dtypes(include=np.number).columns.tolist()
+        categorical_features = X.select_dtypes(exclude=np.number).columns.tolist()
+        X_train_processed, X_test_processed, feature_names, _ = preprocess_data(X_train, X_test, numerical_features, categorical_features)
+        if st.session_state.problem_type in ["Binary Classification", "Multi-Class"] and y is not None:
+            y_train = st.session_state.le.transform(y_train) if st.session_state.le else y_train
+            y_test = st.session_state.le.transform(y_test) if st.session_state.le else y_test
+            if st.session_state.problem_type == "Multi-Class":
+                y_train = tf.keras.utils.to_categorical(y_train)
+                y_test = tf.keras.utils.to_categorical(y_test)
+    else:
+        X_train, X_test, y_train, y_test = train_test_split(st.session_state.images, st.session_state.labels, test_size=0.2, random_state=42) if st.session_state.problem_type == "Image Classification" else (st.session_state.images, st.session_state.images.copy(), None, None)
+        X_train_processed, X_test_processed = X_train, X_test
+        feature_names = ["image_features"]
+        if st.session_state.problem_type == "Image Classification":
+            le = None
+
     model = st.session_state.model
-    preprocessor = st.session_state.preprocessor
-    features = st.session_state.features
-    target = st.session_state.target
     problem_type = st.session_state.problem_type
-    le = st.session_state.le
-
-    X = df[features]
-    y = df[target] if problem_type != "Clustering" else None
-    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) if problem_type != "Clustering" else (X, X.copy(), None, None)
-
-    # FIXED FEATURE SELECTION
-    numerical_features = X.select_dtypes(include=np.number).columns.tolist()
-    categorical_features = X.select_dtypes(exclude=np.number).columns.tolist()
-    X_train_processed, X_test_processed, feature_names, _ = preprocess_data(X_train, X_test, numerical_features, categorical_features)
-
-    if problem_type in ["Binary Classification", "Multi-Class"] and y is not None:
-        y_train = le.transform(y_train) if le else y_train
-        y_test = le.transform(y_test) if le else y_test
-        if problem_type == "Multi-Class":
-            y_train = tf.keras.utils.to_categorical(y_train)
-            y_test = tf.keras.utils.to_categorical(y_test)
+    encoder = st.session_state.get('encoder', None)
 
     # Validation
     st.subheader("Model Validation")
-    metrics, y_pred = evaluate_model(model, X_test_processed, y_test, problem_type, le)
+    if problem_type == "Compression":
+        metrics, y_pred, compressed_data = evaluate_model(model, X_test_processed, None, problem_type, None, encoder)
+    else:
+        metrics, y_pred = evaluate_model(model, X_test_processed, y_test, problem_type, st.session_state.le if data_type == "Tabular" else None)
+
     col1, col2 = st.columns(2)
     with col1:
         for metric, value in metrics.items():
             st.metric(metric, f"{value:.4f}" if isinstance(value, float) else value)
-
     with col2:
         if problem_type == "Regression":
             fig = px.scatter(x=y_test, y=y_pred, labels={"x": "Actual", "y": "Predicted"}, title="Predicted vs Actual")
             st.plotly_chart(fig)
-        elif problem_type == "Binary Classification":
-            y_pred_proba = model.predict_proba(X_test_processed)[:, 1] if hasattr(model, 'predict_proba') else y_pred # Handle cases where predict_proba isn't available
-            fpr, tpr, _ = roc_curve(y_test, y_pred_proba)
-            roc_auc = auc(fpr, tpr)
-            fig = px.area(x=fpr, y=tpr, title=f"ROC Curve (AUC = {roc_auc:.2f})",
-                         labels={"x": "False Positive Rate", "y": "True Positive Rate"})
-            st.plotly_chart(fig)
+        elif problem_type in ["Binary Classification", "Image Classification"]:
+            if problem_type == "Binary Classification":
+                y_pred_proba = model.predict_proba(X_test_processed)[:, 1] if hasattr(model, 'predict_proba') else y_pred
+                fpr, tpr, _ = roc_curve(y_test, y_pred_proba)
+                roc_auc = auc(fpr, tpr)
+                fig = px.area(x=fpr, y=tpr, title=f"ROC Curve (AUC = {roc_auc:.2f})", labels={"x": "False Positive Rate", "y": "True Positive Rate"})
+                st.plotly_chart(fig)
+            else:
+                cm = np.zeros((len(st.session_state.class_names), len(st.session_state.class_names)))
+                for i, j in zip(y_test, y_pred):
+                    cm[i, j] += 1
+                fig = px.imshow(cm, title="Confusion Matrix", labels={"x": "Predicted", "y": "Actual"})
+                st.plotly_chart(fig)
+                report = classification_report(y_test, y_pred, target_names=st.session_state.class_names, zero_division=0)
+                st.text("Classification Report:\n" + report)
         elif problem_type == "Multi-Class":
-            y_pred_classes = np.argmax(model.predict(X_test_processed), axis=1) if isinstance(model, keras.Model) else model.predict(X_test_processed)
+            y_pred_classes = np.argmax(model.predict(X_test_processed), axis=1)
             y_test_classes = np.argmax(y_test, axis=1)
             cm = np.zeros((y_train.shape[1], y_train.shape[1]))
             for i, j in zip(y_test_classes, y_pred_classes):
                 cm[i, j] += 1
             fig = px.imshow(cm, title="Confusion Matrix", labels={"x": "Predicted", "y": "Actual"})
             st.plotly_chart(fig)
-
-            # Display Classification Report
-            report = classification_report(y_test_classes, y_pred_classes, target_names=le.classes_ if le else [str(i) for i in range(y_train.shape[1])], zero_division=0)
+            report = classification_report(y_test_classes, y_pred_classes, target_names=st.session_state.le.classes_ if st.session_state.le else [str(i) for i in range(y_train.shape[1])], zero_division=0)
             st.text("Classification Report:\n" + report)
         elif problem_type == "Clustering":
             labels = y_pred
-            fig = px.scatter(x=X_test_processed[:, 0], y=X_test_processed[:, 1], color=labels, title="Cluster Visualization")
+            fig = px.scatter(x=X_test_processed[:, 0] if X_test_processed.shape[-1] == 1 else X_test_processed.reshape(X_test_processed.shape[0], -1)[:, 0], 
+                            y=X_test_processed[:, 1] if X_test_processed.shape[-1] == 1 else X_test_processed.reshape(X_test_processed.shape[0], -1)[:, 1], 
+                            color=labels, title="Cluster Visualization")
+            st.plotly_chart(fig)
+        elif problem_type == "Compression":
+            st.image([X_test_processed[0], y_pred[0]], caption=["Original", "Reconstructed"], width=200) if data_type == "Image" else None
+            fig = px.scatter(x=X_test_processed.flatten()[:1000], y=y_pred.flatten()[:1000], labels={"x": "Original", "y": "Reconstructed"}, title="Original vs Reconstructed (First 1000 Values)")
             st.plotly_chart(fig)
 
     # Dimensionality Reduction
     st.subheader("Dimensionality Reduction")
-    method = st.selectbox("Method", ["PCA", "SVD", "t-SNE", "UMAP"])
-    n_components = st.slider("Components", 2, min(X_train_processed.shape[1], 10), 2)
-
-    if method == "PCA":
-        reducer = PCA(n_components=n_components)
-        X_reduced = reducer.fit_transform(X_train_processed)
-        fig = px.bar(x=range(n_components), y=reducer.explained_variance_ratio_, title="Explained Variance Ratio")
-        st.plotly_chart(fig)
-    elif method == "SVD":
-        reducer = TruncatedSVD(n_components=n_components)
-        X_reduced = reducer.fit_transform(X_train_processed)
-        fig = px.bar(x=range(n_components), y=reducer.explained_variance_ratio_, title="Explained Variance Ratio")
-        st.plotly_chart(fig)
-    elif method == "t-SNE":
-        X_reduced = TSNE(n_components=n_components, random_state=42).fit_transform(X_train_processed)
-    elif method == "UMAP":
-        X_reduced = umap.UMAP(n_components=n_components, random_state=42).fit_transform(X_train_processed)
+    methods = ["PCA", "SVD", "t-SNE", "UMAP"]
+    if problem_type == "Compression" and 'encoder' in st.session_state:
+        methods.append("Autoencoder")
+    method = st.selectbox("Method", methods)
+    n_components = st.slider("Components", 2, min(10 if data_type == "Tabular" else X_train_processed.shape[1], 10), 2)
+
+    if method == "Autoencoder" and 'encoder' in st.session_state:
+        X_reduced = st.session_state.encoder.predict(X_train_processed)
+        if X_reduced.shape[1] < n_components:
+            st.warning(f"Autoencoder encoding dimension is {X_reduced.shape[1]}, using that instead of {n_components}.")
+            n_components = X_reduced.shape[1]
+    else:
+        X_flat = X_train_processed if data_type == "Tabular" else X_train_processed.reshape(X_train_processed.shape[0], -1)
+        if method == "PCA":
+            reducer = PCA(n_components=n_components)
+            X_reduced = reducer.fit_transform(X_flat)
+            fig = px.bar(x=range(n_components), y=reducer.explained_variance_ratio_, title="Explained Variance Ratio")
+            st.plotly_chart(fig)
+        elif method == "SVD":
+            reducer = TruncatedSVD(n_components=n_components)
+            X_reduced = reducer.fit_transform(X_flat)
+            fig = px.bar(x=range(n_components), y=reducer.explained_variance_ratio_, title="Explained Variance Ratio")
+            st.plotly_chart(fig)
+        elif method == "t-SNE":
+            with st.spinner("Running t-SNE..."):
+                X_reduced = TSNE(n_components=n_components, random_state=42).fit_transform(X_flat)
+        elif method == "UMAP":
+            with st.spinner("Running UMAP..."):
+                X_reduced = umap.UMAP(n_components=n_components, random_state=42).fit_transform(X_flat)
 
     if n_components >= 2:
-        fig = px.scatter(x=X_reduced[:, 0], y=X_reduced[:, 1], color=y_train if problem_type != "Clustering" else y_pred,
-                        title=f"{method} Visualization")
+        if n_components == 2:
+            fig = px.scatter(x=X_reduced[:, 0], y=X_reduced[:, 1], color=y_train if problem_type not in ["Clustering", "Compression"] else y_pred,
+                            title=f"{method} Visualization")
+        elif n_components == 3:
+            fig = px.scatter_3d(x=X_reduced[:, 0], y=X_reduced[:, 1], z=X_reduced[:, 2], color=y_train if problem_type not in ["Clustering", "Compression"] else y_pred,
+                               title=f"{method} Visualization")
         st.plotly_chart(fig)
 
     # Interpretability
-    st.subheader("Interpretability")
-    try:
-        explainer = shap.KernelExplainer(model.predict, X_test_processed[:50]) if isinstance(model, keras.Model) else shap.Explainer(model, X_test_processed)
-        shap_values = explainer.shap_values(X_test_processed[:50])
-
-        if problem_type == "Regression":
-            shap_fig, ax = plt.subplots()
-            shap.summary_plot(shap_values, X_test_processed[:50], feature_names=feature_names, show=False)
-            st.pyplot(shap_fig)
-        elif problem_type in ["Binary Classification", "Multi-Class"]:
-            class_names = le.classes_ if le else [str(i) for i in range(y_train.shape[1])]
-            for i in range(len(class_names)):
+    if problem_type not in ["Compression", "Clustering"]:
+        st.subheader("Interpretability")
+        try:
+            X_flat = X_test_processed if data_type == "Tabular" else X_test_processed.reshape(X_test_processed.shape[0], -1)
+            if isinstance(model, keras.Model):
+                explainer = shap.DeepExplainer(model, X_train_processed[:50] if data_type == "Tabular" else X_train_processed[:50])
+                shap_values = explainer.shap_values(X_flat[:50])
+            else:
+                explainer = shap.Explainer(model, X_flat)
+                shap_values = explainer.shap_values(X_flat[:50])
+            if problem_type == "Regression":
                 shap_fig, ax = plt.subplots()
-                shap.summary_plot(shap_values[i], X_test_processed[:50], feature_names=feature_names, class_names=class_names, show=False)
+                shap.summary_plot(shap_values, X_flat[:50], feature_names=feature_names, show=False)
                 st.pyplot(shap_fig)
-        else:
-            st.write("SHAP plots are not directly applicable to Clustering problems.")
-    except Exception as e:
-        st.error(f"Error generating SHAP plot: {e}")
+            elif problem_type in ["Binary Classification", "Multi-Class", "Image Classification"]:
+                class_names = st.session_state.le.classes_ if st.session_state.le and data_type == "Tabular" else st.session_state.class_names if problem_type == "Image Classification" else [str(i) for i in range(y_train.shape[1])]
+                for i in range(min(len(class_names), len(shap_values))):
+                    shap_fig, ax = plt.subplots()
+                    shap.summary_plot(shap_values[i] if isinstance(shap_values, list) else shap_values, X_flat[:50], feature_names=feature_names, class_names=class_names, show=False)
+                    st.pyplot(shap_fig)
+        except Exception as e:
+            st.error(f"Error generating SHAP plot: {e}")
 
 # Custom CSS
 st.markdown("""