Update app.py

app.py

@@ -1,733 +1,733 @@
-import gradio as gr
-import os
-from PIL import Image
-import numpy as np
-import pickle
-import io
-import sys
-import torch
-import subprocess
-import h5py
-from sklearn.metrics import confusion_matrix
-import matplotlib.pyplot as plt
-import pandas as pd
-from sklearn.metrics import f1_score
-import seaborn as sns
-
-#################### BEAM PREDICTION #########################}
-def beam_prediction_task(data_percentage, task_complexity, theme='Dark'):
-    # Folder naming convention based on input_type, data_percentage, and task_complexity
-    raw_folder = f"images/raw_{data_percentage/100:.1f}_{task_complexity}"
-    embeddings_folder = f"images/embedding_{data_percentage/100:.1f}_{task_complexity}"
-
-    # Process raw confusion matrix
-    raw_cm = compute_average_confusion_matrix(raw_folder)
-    if raw_cm is not None:
-        raw_cm_path = os.path.join(raw_folder, "confusion_matrix_raw.png")
-        plot_confusion_matrix_beamPred(raw_cm, 
-                                       classes=np.arange(raw_cm.shape[0]), 
-                                       title=f"Confusion Matrix (Raw Channels)\n{data_percentage}% data, {task_complexity} beams", 
-                                       save_path=raw_cm_path, 
-                                       theme=theme)
-        raw_img = Image.open(raw_cm_path)
-    else:
-        raw_img = None
-
-    # Process embeddings confusion matrix
-    embeddings_cm = compute_average_confusion_matrix(embeddings_folder)
-    if embeddings_cm is not None:
-        embeddings_cm_path = os.path.join(embeddings_folder, "confusion_matrix_embeddings.png")
-        plot_confusion_matrix_beamPred(embeddings_cm, 
-                                       classes=np.arange(embeddings_cm.shape[0]), 
-                                       title=f"Confusion Matrix (LWM Embeddings)\n{data_percentage}% data, {task_complexity} beams", 
-                                       save_path=embeddings_cm_path, 
-                                       theme=theme)
-        embeddings_img = Image.open(embeddings_cm_path)
-    else:
-        embeddings_img = None
-
-    return raw_img, embeddings_img
-
-# Function to compute the F1-score based on the confusion matrix
-def compute_f1_score(cm):
-    # Compute precision and recall
-    TP = np.diag(cm)
-    FP = np.sum(cm, axis=0) - TP
-    FN = np.sum(cm, axis=1) - TP
-    
-    precision = TP / (TP + FP)
-    recall = TP / (TP + FN)
-    
-    # Handle division by zero in precision or recall
-    precision = np.nan_to_num(precision)
-    recall = np.nan_to_num(recall)
-    
-    # Compute F1 score
-    f1 = 2 * (precision * recall) / (precision + recall)
-    f1 = np.nan_to_num(f1)  # Replace NaN with 0
-    return np.mean(f1)  # Return the mean F1-score across all classes
-
-def plot_confusion_matrix_beamPred(cm, classes, title, save_path, theme='Dark'):
-    # Compute the average F1-score
-    avg_f1 = compute_f1_score(cm)
-
-    # Choose the color scheme based on the user's mode
-    if theme == 'Dark':
-        plt.style.use('dark_background')  # Use dark mode styling
-        #text_color = 'white'
-        text_color = 'gray'
-        #cmap = 'cividis'  # Dark-mode-friendly colormap
-        cmap = 'coolwarm'
-    else:
-        plt.style.use('default')  # Use default (light) mode styling
-        #text_color = 'black'
-        text_color = 'gray'
-        cmap = 'Blues'  # Light-mode-friendly colormap
-
-    plt.figure(figsize=(10, 10))
-
-    # Plot the confusion matrix with a colormap compatible for the mode
-    ax = sns.heatmap(cm, cmap=cmap, cbar=True)
-
-    cbar = ax.collections[0].colorbar  
-    cbar.ax.yaxis.set_tick_params(color=text_color) 
-    cbar.ax.yaxis.set_tick_params(labelcolor=text_color)
-
-    # Add F1-score to the title
-    plt.title(f"{title}\nF1 Score: {avg_f1:.3f}", color=text_color, fontsize=23)
-    
-    tick_marks = np.arange(len(classes))
-    plt.xticks(tick_marks, classes, color=text_color, fontsize=14)  # Adjust text color based on the mode
-    plt.yticks(tick_marks, classes, color=text_color, fontsize=14)  # Adjust text color based on the mode
-
-    plt.ylabel('True label', color=text_color, fontsize=20)
-    plt.xlabel('Predicted label', color=text_color, fontsize=20)
-    plt.tight_layout()
-
-    plt.savefig(save_path, transparent=True)  # Transparent to blend with the site background
-    plt.close()
-
-    # Return the saved image
-    return Image.open(save_path)
-
-def compute_average_confusion_matrix(folder):
-    confusion_matrices = []
-    max_num_labels = 0
-
-    # First pass to determine the maximum number of labels
-    for file in os.listdir(folder):
-        if file.endswith(".csv"):
-            data = pd.read_csv(os.path.join(folder, file))
-            num_labels = len(np.unique(data["Target"]))
-            max_num_labels = max(max_num_labels, num_labels)
-
-    # Second pass to calculate the confusion matrices and pad if necessary
-    for file in os.listdir(folder):
-        if file.endswith(".csv"):
-            data = pd.read_csv(os.path.join(folder, file))
-            y_true = data["Target"]
-            y_pred = data["Top-1 Prediction"]
-            num_labels = len(np.unique(y_true))
-            
-            # Compute confusion matrix
-            cm = confusion_matrix(y_true, y_pred, labels=np.arange(max_num_labels))
-
-            # If the confusion matrix is smaller, pad it to match the largest size
-            if cm.shape[0] < max_num_labels:
-                padded_cm = np.zeros((max_num_labels, max_num_labels))
-                padded_cm[:cm.shape[0], :cm.shape[1]] = cm
-                confusion_matrices.append(padded_cm)
-            else:
-                confusion_matrices.append(cm)
-
-    if confusion_matrices:
-        avg_cm = np.mean(confusion_matrices, axis=0)
-        return avg_cm
-    else:
-        return None
-
-########################## LOS/NLOS CLASSIFICATION #############################3
-
-
-# Paths to the predefined images folder
-LOS_PATH = "images_LoS"
-
-# Define the percentage values
-percentage_values_los = np.linspace(0.001, 1, 20) * 100 #np.linspace(0.05, 1, 20) * 100 # np.linspace(0.001, 1, 20) * 100  # 20 percentage values
-
-from sklearn.metrics import f1_score
-import seaborn as sns
-
-# Function to compute confusion matrix, F1-score and plot it with dark mode style
-def plot_confusion_matrix_from_csv(csv_file_path, title, save_path, light_mode=False):
-    # Load CSV file
-    data = pd.read_csv(csv_file_path)
-    
-    # Extract ground truth and predictions
-    y_true = data['Target']
-    y_pred = data['Top-1 Prediction']
-    
-    # Compute confusion matrix
-    cm = confusion_matrix(y_true, y_pred)
-    
-    # Compute F1-score
-    f1 = f1_score(y_true, y_pred, average='macro')  # Macro-average F1-score
-
-    # Set styling based on light or dark mode
-    if light_mode:
-        plt.style.use('default')  # Light mode styling
-        text_color = 'black'
-        cmap = 'Blues'  # Light-mode-friendly colormap
-    else:
-        plt.style.use('dark_background')  # Dark mode styling
-        text_color = 'gray'
-        #cmap = 'magma'  # Dark-mode-friendly colormap
-        cmap = 'coolwarm'
-    
-    plt.figure(figsize=(5, 5))
-
-    # Plot the confusion matrix with the chosen colormap
-    sns.heatmap(cm, annot=True, fmt="d", cmap=cmap, cbar=False, annot_kws={"size": 12}, linewidths=0.5, linecolor='white')
-
-    # Add F1-score to the title
-    plt.title(f"{title}\nF1 Score: {f1:.3f}", color=text_color, fontsize=18)
-
-    # Customize tick labels for light/dark mode
-    plt.xticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
-    plt.yticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
-
-    plt.ylabel('True label', color=text_color, fontsize=14)
-    plt.xlabel('Predicted label', color=text_color, fontsize=14)
-    plt.tight_layout()
-
-    # Save the plot as an image
-    plt.savefig(save_path, transparent=True)  # Use transparent to blend with the website
-    plt.close()
-
-    # Return the saved image
-    return Image.open(save_path)
-
-# Function to load confusion matrix based on percentage and input_type
-def display_confusion_matrices_los(percentage):
-    #percentage = percentage_values_los[percentage_idx]
-    
-    # Construct folder names
-    raw_folder = os.path.join(LOS_PATH, f"raw_{percentage/100:.3f}_los_noTraining")
-    embeddings_folder = os.path.join(LOS_PATH, f"embedding_{percentage/100:.3f}_los_noTraining")
-    
-    # Process raw confusion matrix
-    raw_csv_file = os.path.join(raw_folder, f"test_predictions_raw_{percentage/100:.3f}_los.csv")
-    raw_cm_img_path = os.path.join(raw_folder, "confusion_matrix_raw.png")
-    raw_img = plot_confusion_matrix_from_csv(raw_csv_file, 
-                                             f"Confusion Matrix (Raw Channels)\n{percentage:.1f}% data", 
-                                             raw_cm_img_path)
-
-    # Process embeddings confusion matrix
-    embeddings_csv_file = os.path.join(embeddings_folder, f"test_predictions_embedding_{percentage/100:.3f}_los.csv")
-    embeddings_cm_img_path = os.path.join(embeddings_folder, "confusion_matrix_embeddings.png")
-    embeddings_img = plot_confusion_matrix_from_csv(embeddings_csv_file, 
-                                                    f"Confusion Matrix (LWM Embeddings)\n{percentage:.1f}% data", 
-                                                    embeddings_cm_img_path)
-
-    return raw_img, embeddings_img
-
-# Main function to handle user choice
-def handle_user_choice(choice, percentage=None, uploaded_file=None, emb_type='CLS Embedding'):
-    if choice == "Use Default Dataset":
-        raw_img, embeddings_img = display_confusion_matrices_los(percentage)
-        return raw_img, embeddings_img, ""  # Return empty string for console output
-    elif choice == "Upload Dataset":
-        if uploaded_file is not None:
-            raw_img, embeddings_img, console_output = process_hdf5_file(uploaded_file, percentage, emb_type)
-            return raw_img, embeddings_img, console_output
-        else:
-            return "Please upload a dataset", "Please upload a dataset", ""  # Return empty string for console output
-    else:
-        return "Invalid choice", "Invalid choice", ""  # Return empty string for console output
-
-# Custom class to capture print output
-class PrintCapture(io.StringIO):
-    def __init__(self):
-        super().__init__()
-        self.output = []
-
-    def write(self, txt):
-        self.output.append(txt)
-        super().write(txt)
-
-    def get_output(self):
-        return ''.join(self.output)
-
-# Function to load and display predefined images based on user selection
-def display_predefined_images(percentage):
-    #percentage = percentage_values_los[percentage_idx]
-    raw_image_path = os.path.join(RAW_PATH, f"percentage_{percentage}_complexity_16.png")
-    embeddings_image_path = os.path.join(EMBEDDINGS_PATH, f"percentage_{percentage}_complexity_16.png")
-    
-    # Check if the images exist
-    if os.path.exists(raw_image_path):
-        raw_image = Image.open(raw_image_path)
-    else:
-        raw_image = create_random_image()  # Use a fallback random image
-    
-    if os.path.exists(embeddings_image_path):
-        embeddings_image = Image.open(embeddings_image_path)
-    else:
-        embeddings_image = create_random_image()  # Use a fallback random image
-
-    return raw_image, embeddings_image
-
-def los_nlos_classification(file, percentage):
-    if file is not None:
-        raw_cm_image, emb_cm_image, console_output = process_hdf5_file(file, percentage)
-        return raw_cm_image, emb_cm_image, console_output  # Returning all three: two images and console output
-    else:
-        raw_image, embeddings_image = display_predefined_images(percentage)
-        return raw_image, embeddings_image, ""  # Return an empty string for console output when no file is uploaded
-
-# Function to create random images for LoS/NLoS classification results
-def create_random_image(size=(300, 300)):
-    random_image = np.random.rand(*size, 3) * 255
-    return Image.fromarray(random_image.astype('uint8'))
-
-import importlib.util
-
-# Function to dynamically load a Python module from a given file path
-def load_module_from_path(module_name, file_path):
-    spec = importlib.util.spec_from_file_location(module_name, file_path)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-# Function to split dataset into training and test sets based on user selection
-def split_dataset(channels, labels, percentage):
-    #percentage = percentage_values_los[percentage_idx] / 100
-    num_samples = channels.shape[0]
-    train_size = int(num_samples * percentage/100)
-    print(f'Number of Training Samples: {train_size}')
-    
-    indices = np.arange(num_samples)
-    np.random.shuffle(indices)
-    
-    train_idx, test_idx = indices[:train_size], indices[train_size:]
-    
-    train_data, test_data = channels[train_idx], channels[test_idx]
-    train_labels, test_labels = labels[train_idx], labels[test_idx]
-    
-    return train_data, test_data, train_labels, test_labels
-
-# Function to calculate Euclidean distance between a point and a centroid
-def euclidean_distance(x, centroid):
-    return np.linalg.norm(x - centroid)
-
-import torch
-
-def classify_based_on_distance(train_data, train_labels, test_data):
-    # Compute the centroids for the two classes
-    centroid_0 = train_data[train_labels == 0].mean(dim=0)  # Use torch.mean
-    centroid_1 = train_data[train_labels == 1].mean(dim=0)  # Use torch.mean
-    
-    predictions = []
-    for test_point in test_data:
-        # Compute Euclidean distance between the test point and each centroid
-        dist_0 = euclidean_distance(test_point, centroid_0)
-        dist_1 = euclidean_distance(test_point, centroid_1)
-        predictions.append(0 if dist_0 < dist_1 else 1)
-    
-    return torch.tensor(predictions)  # Return predictions as a PyTorch tensor
-
-def plot_confusion_matrix(y_true, y_pred, title, light_mode=False):
-    cm = confusion_matrix(y_true, y_pred)
-    
-    # Calculate F1 Score
-    f1 = f1_score(y_true, y_pred, average='weighted')
-
-    #plt.style.use('dark_background')
-    
-    # Set styling based on light or dark mode
-    if light_mode:
-        plt.style.use('default')  # Light mode styling
-        text_color = 'black'
-        cmap = 'Blues'  # Light-mode-friendly colormap
-    else:
-        plt.style.use('dark_background')  # Dark mode styling
-        text_color = 'gray'
-        #cmap = 'magma'  # Dark-mode-friendly colormap
-        cmap = 'coolwarm'
-        
-    plt.figure(figsize=(5, 5))
-    
-    # Plot the confusion matrix with a dark-mode compatible colormap
-    sns.heatmap(cm, annot=True, fmt="d", cmap=cmap, cbar=False, annot_kws={"size": 12}, linewidths=0.5, linecolor='white')
-    
-    # Add F1-score to the title
-    plt.title(f"{title}\nF1 Score: {f1:.3f}", color=text_color, fontsize=18)
-    
-    # Customize tick labels for dark mode
-    plt.xticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
-    plt.yticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
-    
-    plt.ylabel('True label', color=text_color, fontsize=14)
-    plt.xlabel('Predicted label', color=text_color, fontsize=14)
-    plt.tight_layout()
-    
-    # Save the plot as an image
-    plt.savefig(f"{title}.png", transparent=True)  # Use transparent to blend with the dark mode website
-    plt.close()
-    
-    # Return the saved image
-    return Image.open(f"{title}.png")
-
-
-def identical_train_test_split(output_emb, output_raw, labels, train_percentage):
-
-    torch.manual_seed(42)
-    
-    N = output_emb.shape[0]
-    indices = torch.randperm(N) 
-    test_split_index = int(N * 0.20)
-    test_indices = indices[:test_split_index]
-    remaining_indices = indices[test_split_index:]
-    train_split_index = int(len(remaining_indices) * train_percentage / 100)
-    print(f'Training Size: {train_split_index} out of remaining {len(remaining_indices)}')
-    print(f'Test Size: {test_split_index}')
-    
-    train_indices = remaining_indices[:train_split_index]
-
-    train_emb = output_emb[train_indices]
-    test_emb = output_emb[test_indices]
-    
-    train_raw = output_raw[train_indices]
-    test_raw = output_raw[test_indices]
-
-    train_labels = labels[train_indices]
-    test_labels = labels[test_indices]
-
-    return train_emb, test_emb, train_raw, test_raw, train_labels, test_labels
-
-# Store the original working directory when the app starts
-original_dir = os.getcwd()
-
-def process_hdf5_file(uploaded_file, percentage, emb_type='CLS Embedding'):
-    capture = PrintCapture()
-    sys.stdout = capture  # Redirect print statements to capture
-    
-    try:
-        model_repo_url = "https://huggingface.co/wi-lab/lwm"
-        model_repo_dir = "./LWM"
-
-        # Step 1: Clone the repository if not already done
-        if not os.path.exists(model_repo_dir):
-            print(f"Cloning model repository from {model_repo_url}...")
-            subprocess.run(["git", "clone", model_repo_url, model_repo_dir], check=True)
-
-        # Step 2: Verify the repository was cloned and change the working directory
-        repo_work_dir = os.path.join(original_dir, model_repo_dir)
-        if os.path.exists(repo_work_dir):
-            os.chdir(repo_work_dir)  # Change the working directory only once
-            print(f"Changed working directory to {os.getcwd()}")
-            #print(f"Directory content: {os.listdir(os.getcwd())}")  # Debugging: Check repo content
-        else:
-            print(f"Directory {repo_work_dir} does not exist.")
-            return
-            
-        # Step 3: Dynamically load lwm_model.py, input_preprocess.py, and inference.py
-        lwm_model_path = os.path.join(os.getcwd(), 'lwm_model.py')
-        input_preprocess_path = os.path.join(os.getcwd(), 'input_preprocess.py')
-        inference_path = os.path.join(os.getcwd(), 'inference.py')
-
-        # Load lwm_model
-        lwm_model = load_module_from_path("lwm_model", lwm_model_path)
-
-        # Load input_preprocess
-        input_preprocess = load_module_from_path("input_preprocess", input_preprocess_path)
-
-        # Load inference
-        inference = load_module_from_path("inference", inference_path)
-
-        # Step 4: Load the model from lwm_model module
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        print(f"Loading the LWM model on {device}...")
-        model = lwm_model.lwm.from_pretrained(device=device).float()
-
-        # Step 5: Load the HDF5 file and extract the channels and labels
-        with h5py.File(uploaded_file.name, 'r') as f:
-            channels = np.array(f['channels']).astype(np.complex64)
-            labels = np.array(f['labels']).astype(np.int32)  
-        print(f"Loaded dataset with {channels.shape[0]} samples.")
-        channels = channels * (10**(-3-np.floor(np.log10(np.abs(np.mean(channels).real)))))
-        
-        # Step 7: Tokenize the data using the tokenizer from input_preprocess
-        preprocessed_chs = input_preprocess.tokenizer(manual_data=channels)
-
-        # Step 7: Perform inference using the functions from inference.py
-        if emb_type == 'Channel Embedding':
-            embedding_type = 'channel_emb'
-        elif emb_type == 'CLS Embedding':
-            embedding_type = 'cls_emb'
-            
-        output_emb = inference.lwm_inference(preprocessed_chs, embedding_type, model, device)
-        output_raw = inference.create_raw_dataset(preprocessed_chs, device)
-
-        print(f"Output Embeddings Shape: {output_emb.shape}")
-        print(f"Output Raw Shape: {output_raw.shape}")
-
-        print(f'percentage_value: {percentage}')
-        train_data_emb, test_data_emb, train_data_raw, test_data_raw, train_labels, test_labels = identical_train_test_split(output_emb.view(len(output_emb),-1),
-                                                                                                                             output_raw.view(len(output_raw),-1),
-                                                                                                                             labels,
-                                                                                                                             percentage)
-        
-        # Step 8: Perform classification using the Euclidean distance for both raw and embeddings
-        #print(f'train_data_emb: {train_data_emb.shape}')
-        #print(f'train_labels: {train_labels.shape}')
-        #print(f'test_data_emb: {test_data_emb.shape}')
-        pred_raw = classify_based_on_distance(train_data_raw, train_labels, test_data_raw)
-        pred_emb = classify_based_on_distance(train_data_emb, train_labels, test_data_emb)
-        
-        # Step 9: Generate confusion matrices for both raw and embeddings
-        raw_cm_image = plot_confusion_matrix(test_labels, pred_raw, title="Confusion Matrix (Raw Channels)")
-        emb_cm_image = plot_confusion_matrix(test_labels, pred_emb, title="Confusion Matrix (Embeddings)")
-
-        return raw_cm_image, emb_cm_image, capture.get_output()
-
-    except Exception as e:
-        return str(e), str(e), capture.get_output()
-
-    finally:
-        # Always return to the original working directory after processing
-        os.chdir(original_dir)
-        sys.stdout = sys.__stdout__  # Reset print statements
-
-######################## Define the Gradio interface ###############################
-js_code = """
-() => {
-  const isDarkMode = window.matchMedia && window.matchMedia('(prefers-color-scheme: dark)').matches;
-  return isDarkMode ? 'dark' : 'light';
-}
-"""
-
-with gr.Blocks(css="""
-    .slider-container {
-        display: inline-block;
-        margin-right: 50px;
-        text-align: center;
-    }
-
-    .explanation-box {
-        font-size: 16px;
-        font-style: italic;
-        color: #4a4a4a;
-        padding: 15px;
-        background-color: #f0f0f0;
-        border-radius: 10px;
-        margin-bottom: 20px;
-    }
-
-    .bold-highlight {
-        font-weight: bold;
-        color: #2c3e50;
-        font-size: 18px;
-        text-align: center;
-        margin-bottom: 20px;
-    }
-
-    #console-output {
-        background-color: #ffffff;  /* Light background for light mode */
-        color: #000000;  /* Dark text color for contrast */
-        padding: 10px;
-        border-radius: 5px;
-    }
-
-    .plot-title {
-        font-weight: bold;
-        color: #2c3e50;
-    }
-""") as demo:
-
-    # Contact Section
-    gr.Markdown("""
-        <div style="text-align: center;">
-            <a target="_blank" href="https://www.wi-lab.net">
-                <img src="https://www.wi-lab.net/wp-content/uploads/2021/08/WI-name.png" alt="Wireless Model" style="height: 30px;">
-            </a>
-            <a target="_blank" href="mailto:lwmwireless@gmail.com" style="margin-left: 10px;">
-                <img src="https://img.shields.io/badge/email-lwmwireless@gmail.com-blue.svg?logo=gmail" alt="Email">
-            </a>
-        </div>
-    """)
-    
-    gr.Markdown("""
-    <div class="bold-highlight">
-        🚀 Explore the pre-trained <b>LWM Model<b> here: 
-        <a target="_blank" href="https://huggingface.co/wi-lab/lwm">https://huggingface.co/wi-lab/lwm</a>
-    </div>
-    """)
-    
-    # Tab for Beam Prediction Task
-    with gr.Tab("Beam Prediction Task"):
-        #gr.Markdown("### Beam Prediction Task")
-        
-        # Explanation section with creative spacing and minimal design
-        gr.Markdown("""
-            <div style="background-color: var(--primary-background); padding: 15px; border-radius: 10px; color: var(--text-primary);">
-              <h3 style="color: var(--text-primary);">📡 <b>Sub-6GHz to mmWave Beam Prediction Task</b></h3>
-              <ul style="padding-left: 20px;">
-                <li><b>🎯 Goal</b>: Predict the strongest <b>mmWave beam</b> from a predefined codebook using Sub-6 GHz channels.</li>
-                <li><b>⚙️ Adjust Settings</b>: Use the sliders to control the training data percentage and task complexity (beam count) to explore model performance.</li>
-                <li><b>🧠 Inferences</b>:
-                  <ul>
-                    <li>🔍 First, the LWM model extracts features.</li>
-                    <li>🤖 Then, the downstream residual 1D-CNN model (500K parameters) makes beam predictions.</li>
-                  </ul>
-                </li>
-                <li><b>🗺️ Dataset</b>: A combination of six scenarios from the DeepMIMO dataset (excluded from LWM pre-training) highlights the model's strong generalization abilities.</li>
-              </ul>
-            </div>
-        """)
-
-        with gr.Row():
-            with gr.Column():
-                data_percentage_slider = gr.Slider(label="Data Percentage for Training", minimum=10, maximum=100, step=10, value=10)
-                task_complexity_dropdown = gr.Dropdown(label="Task Complexity (Number of Beams)", choices=[16, 32, 64, 128, 256], value=16)
-                #theme_dropdown = gr.Dropdown(label="Select Theme", choices=['Light', 'Dark'], value='Light')
-
-        with gr.Row():
-            raw_img_bp = gr.Image(label="Raw Channels", type="pil", width=300, height=500)
-            embeddings_img_bp = gr.Image(label="Embeddings", type="pil", width=300, height=500)
-
-        theme_dropdown = 'Dark'
-        # Update the confusion matrices whenever sliders change
-        data_percentage_slider.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
-        task_complexity_dropdown.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
-        #theme_dropdown.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown, theme_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
-        
-        # Add a conclusion section at the bottom
-        gr.Markdown("""
-            <div class="explanation-box">
-                The LWM embeddings demonstrate remarkable generalization capabilities, enabling impressive performance even with minimal training samples. This highlights their ability to effectively handle diverse tasks with limited data.
-            </div>
-        """)
-
-    # Separate Tab for LoS/NLoS Classification Task
-    with gr.Tab("LoS/NLoS Classification Task"):
-        #gr.Markdown("### LoS/NLoS Classification Task")
-
-        # Explanation section with creative spacing
-        gr.Markdown("""
-            <div style="background-color: var(--primary-background); padding: 15px; border-radius: 10px; color: var(--text-primary);">
-              <h3 style="color: var(--text-primary);">🔍 <b>LoS/NLoS Classification Task</b></h3>
-              <ul style="padding-left: 20px;">
-                <li><b>🎯 Goal</b>: Classify whether a channel is <b>LoS</b> (Line-of-Sight) or <b>NLoS</b> (Non-Line-of-Sight) with very small LWM CLS embeddings.</li>
-                <li><b>📂 Dataset</b>: Use the default dataset (a combination of six scenarios from the DeepMIMO dataset) or upload your own dataset in <b>H5</b> format.</li>
-                <li><b>💡 Custom Dataset Requirements:</b>
-                  <ul>
-                    <li>📡 <b>channels</b> array: Shape (N,32,32), rows: 32 antennas at BS, columns: 32 subcarriers</li>
-                    <li>🏷️ <b>labels</b> array: Binary LoS/NLoS values (1/0)</li>
-                  </ul>
-                </li>
-                <li><b>🔗 Tip 1</b>: Instructions for organizing your dataset are available at the bottom of the page.</li>
-                <li><b>🔗 Tip 2</b>: As the computations and inference are performed on HuggingFace CPUs, please use small datasets for faster demo experience (say <400 samples). Clone the model from <a href="https://huggingface.co/wi-lab/lwm" target="_blank">here</a> and use any number of samples locally.</li>
-                <li><b>🔗 Tip 3</b>: Your dataset will be normalized automatically based on outdoor environments. </li>
-                <li><b>💼 No Downstream Model</b>: Instead of a complex downstream model, we classify each sample based on its distance to the centroid of training samples from each class (LoS/NLoS).</il>
-              </ul>
-            </div>
-        """)
-        # Radio button for user choice: predefined data or upload dataset
-        choice_radio = gr.Radio(choices=["Use Default Dataset", "Upload Dataset"], 
-                                label="Choose how to proceed", 
-                                value="Use Default Dataset")
-        
-        percentage_slider_los = gr.Slider(minimum=float(percentage_values_los[0]), 
-                                          maximum=float(percentage_values_los[-1]), 
-                                          step=float((percentage_values_los[-1] - percentage_values_los[0]) / (len(percentage_values_los) - 1)), 
-                                          value=float(percentage_values_los[0]), 
-                                          label="Percentage of Data for Training", 
-                                          interactive=True)
-        
-        # File uploader for dataset (only visible if user chooses to upload a dataset)
-        file_input = gr.File(label="Upload HDF5 Dataset", file_types=[".h5"], visible=False)
-        
-        # Dropdown for embedding type, also only visible when "Upload Dataset" is selected
-        emb_type = gr.Dropdown(choices=["Channel Embedding", "CLS Embedding"], 
-                               value="CLS Embedding", 
-                               label="Embedding Type", 
-                               visible=False)
-        
-        # Confusion matrices display
-        with gr.Row():
-            raw_img_los = gr.Image(label="Raw Channels", type="pil", width=300, height=300)
-            embeddings_img_los = gr.Image(label="Embeddings", type="pil", width=300, height=300)
-            output_textbox = gr.Textbox(label="Console Output", lines=10, elem_id="console-output")
-        
-        # Update the visibility of file_input and emb_type based on user choice
-        def toggle_file_input_and_emb_type(choice):
-            visible = (choice == "Upload Dataset")
-            return gr.update(visible=visible), gr.update(visible=visible)
-        
-        # Change visibility of file input and embedding type dropdown based on choice
-        choice_radio.change(fn=toggle_file_input_and_emb_type, 
-                            inputs=[choice_radio], 
-                            outputs=[file_input, emb_type])
-        
-        # When percentage slider changes (for predefined data)
-        percentage_slider_los.change(fn=handle_user_choice, 
-                                     inputs=[choice_radio, percentage_slider_los, file_input, emb_type], 
-                                     outputs=[raw_img_los, embeddings_img_los, output_textbox])
-        
-        # Layout for the UI - this part does NOT need .render(), Gradio will render these automatically
-        with gr.Row():
-            file_input  # No need for .render()
-            emb_type    # No need for .render()
-
-        # Add a conclusion section at the bottom
-        gr.Markdown("""
-            <div class="explanation-box">
-                Despite their compact size (1/32 of the raw channels), LWM CLS embeddings capture rich, holistic information about the channels. This makes them exceptionally well-suited for tasks like LoS/NLoS classification, especially when working with very limited data.
-            </div>
-        """)
-
-        gr.Markdown("""
-            <div class="explanation-box">
-                To create a custom dataset, you'll need to structure your data with 32x32 channel matrices, where the rows correspond to antennas at the base station and the columns represent subcarriers. Here’s how to organize and store the channels and labels in an H5 file format for the demo:
-            </div>
-        
-            ```python
-            # How to pack channels and labels in a h5 file format as a custom dataset for the demo:
-            import h5py
-            with h5py.File('dataset.h5', 'w') as hdf:
-                hdf.create_dataset('channels', data=channels)
-                hdf.create_dataset('labels', data=labels)
-            print("Dataset saved!")
-        """)
-        gr.Markdown("""
-            <div class="explanation-box">
-                To use your preferred DeepMIMO scenarios for the custom dataset, please 
-                <a href="https://huggingface.co/wi-lab/lwm" target="_blank">clone the model and datasets</a> 
-                and follow the instructions below:
-            </div>
-
-            
-            ```python
-            from input_preprocess import DeepMIMO_data_gen deepmimo_data_cleaning label_gen # Import required modules from the model repository
-            import numpy as np
-            scenario_names = np.array([
-                "city_18_denver", "city_15_indianapolis", "city_19_oklahoma", 
-                "city_12_fortworth", "city_11_santaclara", "city_7_sandiego"
-            ])
-            scenario_name = scenario_names[0] # Select the scenario by choosing its index.
-            deepmimo_data = DeepMIMO_data_gen(scenario_name) # Generates ray-traced wireless channels for the selected scenario.
-            cleaned_deepmimo_data = deepmimo_data_cleaning(deepmimo_data) # Filters out users with no direct path to the base station (i.e., users with zero-valued channels).
-            
-            channels = np.squeeze(np.array(cleaned_deepmimo_data), axis=1) # The "channels" array is now prepared for packing into the custom dataset in H5 format.
-            labels = label_gen('LoS/NLoS Classification', deepmimo_data, scenario_name) # Generates labels for each user, classifying them as Line-of-Sight (LoS) or Non-Line-of-Sight (NLoS), and prepares the "labels" array for inclusion in the custom dataset H5 file.
-            ```
-        """)
-    #with gr.Tab("LWM Model and Framework"):    
-    #    gr.Image("images/lwm_model_v2.png")
-    #    gr.Markdown("This figure depicts the offline pre-training and online embedding generation process for LWM. The channel is divided into fixed-size patches, which are linearly embedded and combined with positional encodings before being passed through a Transformer encoder. During self-supervised pre-training, some embeddings are masked, and LWM leverages self-attention to extract deep features, allowing the decoder to reconstruct the masked values. For downstream tasks, the generated LWM embeddings enhance performance.")
-
-# Launch the app
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+import os
+from PIL import Image
+import numpy as np
+import pickle
+import io
+import sys
+import torch
+import subprocess
+import h5py
+from sklearn.metrics import confusion_matrix
+import matplotlib.pyplot as plt
+import pandas as pd
+from sklearn.metrics import f1_score
+import seaborn as sns
+
+#################### BEAM PREDICTION #########################}
+def beam_prediction_task(data_percentage, task_complexity, theme='Dark'):
+    # Folder naming convention based on input_type, data_percentage, and task_complexity
+    raw_folder = f"images/raw_{data_percentage/100:.1f}_{task_complexity}"
+    embeddings_folder = f"images/embedding_{data_percentage/100:.1f}_{task_complexity}"
+
+    # Process raw confusion matrix
+    raw_cm = compute_average_confusion_matrix(raw_folder)
+    if raw_cm is not None:
+        raw_cm_path = os.path.join(raw_folder, "confusion_matrix_raw.png")
+        plot_confusion_matrix_beamPred(raw_cm, 
+                                       classes=np.arange(raw_cm.shape[0]), 
+                                       title=f"Confusion Matrix (Raw Channels)\n{data_percentage}% data, {task_complexity} beams", 
+                                       save_path=raw_cm_path, 
+                                       theme=theme)
+        raw_img = Image.open(raw_cm_path)
+    else:
+        raw_img = None
+
+    # Process embeddings confusion matrix
+    embeddings_cm = compute_average_confusion_matrix(embeddings_folder)
+    if embeddings_cm is not None:
+        embeddings_cm_path = os.path.join(embeddings_folder, "confusion_matrix_embeddings.png")
+        plot_confusion_matrix_beamPred(embeddings_cm, 
+                                       classes=np.arange(embeddings_cm.shape[0]), 
+                                       title=f"Confusion Matrix (LWM Embeddings)\n{data_percentage}% data, {task_complexity} beams", 
+                                       save_path=embeddings_cm_path, 
+                                       theme=theme)
+        embeddings_img = Image.open(embeddings_cm_path)
+    else:
+        embeddings_img = None
+
+    return raw_img, embeddings_img
+
+# Function to compute the F1-score based on the confusion matrix
+def compute_f1_score(cm):
+    # Compute precision and recall
+    TP = np.diag(cm)
+    FP = np.sum(cm, axis=0) - TP
+    FN = np.sum(cm, axis=1) - TP
+    
+    precision = TP / (TP + FP)
+    recall = TP / (TP + FN)
+    
+    # Handle division by zero in precision or recall
+    precision = np.nan_to_num(precision)
+    recall = np.nan_to_num(recall)
+    
+    # Compute F1 score
+    f1 = 2 * (precision * recall) / (precision + recall)
+    f1 = np.nan_to_num(f1)  # Replace NaN with 0
+    return np.mean(f1)  # Return the mean F1-score across all classes
+
+def plot_confusion_matrix_beamPred(cm, classes, title, save_path, theme='Dark'):
+    # Compute the average F1-score
+    avg_f1 = compute_f1_score(cm)
+
+    # Choose the color scheme based on the user's mode
+    if theme == 'Dark':
+        plt.style.use('dark_background')  # Use dark mode styling
+        #text_color = 'white'
+        text_color = 'gray'
+        #cmap = 'cividis'  # Dark-mode-friendly colormap
+        cmap = 'coolwarm'
+    else:
+        plt.style.use('default')  # Use default (light) mode styling
+        #text_color = 'black'
+        text_color = 'gray'
+        cmap = 'Blues'  # Light-mode-friendly colormap
+
+    plt.figure(figsize=(10, 10))
+
+    # Plot the confusion matrix with a colormap compatible for the mode
+    ax = sns.heatmap(cm, cmap=cmap, cbar=True)
+
+    cbar = ax.collections[0].colorbar  
+    cbar.ax.yaxis.set_tick_params(color=text_color) 
+    cbar.ax.yaxis.set_tick_params(labelcolor=text_color)
+
+    # Add F1-score to the title
+    plt.title(f"{title}\nF1 Score: {avg_f1:.3f}", color=text_color, fontsize=23)
+    
+    tick_marks = np.arange(len(classes))
+    plt.xticks(tick_marks, classes, color=text_color, fontsize=14)  # Adjust text color based on the mode
+    plt.yticks(tick_marks, classes, color=text_color, fontsize=14)  # Adjust text color based on the mode
+
+    plt.ylabel('True label', color=text_color, fontsize=20)
+    plt.xlabel('Predicted label', color=text_color, fontsize=20)
+    plt.tight_layout()
+
+    plt.savefig(save_path, transparent=True)  # Transparent to blend with the site background
+    plt.close()
+
+    # Return the saved image
+    return Image.open(save_path)
+
+def compute_average_confusion_matrix(folder):
+    confusion_matrices = []
+    max_num_labels = 0
+
+    # First pass to determine the maximum number of labels
+    for file in os.listdir(folder):
+        if file.endswith(".csv"):
+            data = pd.read_csv(os.path.join(folder, file))
+            num_labels = len(np.unique(data["Target"]))
+            max_num_labels = max(max_num_labels, num_labels)
+
+    # Second pass to calculate the confusion matrices and pad if necessary
+    for file in os.listdir(folder):
+        if file.endswith(".csv"):
+            data = pd.read_csv(os.path.join(folder, file))
+            y_true = data["Target"]
+            y_pred = data["Top-1 Prediction"]
+            num_labels = len(np.unique(y_true))
+            
+            # Compute confusion matrix
+            cm = confusion_matrix(y_true, y_pred, labels=np.arange(max_num_labels))
+
+            # If the confusion matrix is smaller, pad it to match the largest size
+            if cm.shape[0] < max_num_labels:
+                padded_cm = np.zeros((max_num_labels, max_num_labels))
+                padded_cm[:cm.shape[0], :cm.shape[1]] = cm
+                confusion_matrices.append(padded_cm)
+            else:
+                confusion_matrices.append(cm)
+
+    if confusion_matrices:
+        avg_cm = np.mean(confusion_matrices, axis=0)
+        return avg_cm
+    else:
+        return None
+
+########################## LOS/NLOS CLASSIFICATION #############################3
+
+
+# Paths to the predefined images folder
+LOS_PATH = "images_LoS"
+
+# Define the percentage values
+percentage_values_los = np.linspace(0.001, 1, 20) * 100 #np.linspace(0.05, 1, 20) * 100 # np.linspace(0.001, 1, 20) * 100  # 20 percentage values
+
+from sklearn.metrics import f1_score
+import seaborn as sns
+
+# Function to compute confusion matrix, F1-score and plot it with dark mode style
+def plot_confusion_matrix_from_csv(csv_file_path, title, save_path, light_mode=False):
+    # Load CSV file
+    data = pd.read_csv(csv_file_path)
+    
+    # Extract ground truth and predictions
+    y_true = data['Target']
+    y_pred = data['Top-1 Prediction']
+    
+    # Compute confusion matrix
+    cm = confusion_matrix(y_true, y_pred)
+    
+    # Compute F1-score
+    f1 = f1_score(y_true, y_pred, average='macro')  # Macro-average F1-score
+
+    # Set styling based on light or dark mode
+    if light_mode:
+        plt.style.use('default')  # Light mode styling
+        text_color = 'black'
+        cmap = 'Blues'  # Light-mode-friendly colormap
+    else:
+        plt.style.use('dark_background')  # Dark mode styling
+        text_color = 'gray'
+        #cmap = 'magma'  # Dark-mode-friendly colormap
+        cmap = 'coolwarm'
+    
+    plt.figure(figsize=(5, 5))
+
+    # Plot the confusion matrix with the chosen colormap
+    sns.heatmap(cm, annot=True, fmt="d", cmap=cmap, cbar=False, annot_kws={"size": 12}, linewidths=0.5, linecolor='white')
+
+    # Add F1-score to the title
+    plt.title(f"{title}\nF1 Score: {f1:.3f}", color=text_color, fontsize=18)
+
+    # Customize tick labels for light/dark mode
+    plt.xticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
+    plt.yticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
+
+    plt.ylabel('True label', color=text_color, fontsize=14)
+    plt.xlabel('Predicted label', color=text_color, fontsize=14)
+    plt.tight_layout()
+
+    # Save the plot as an image
+    plt.savefig(save_path, transparent=True)  # Use transparent to blend with the website
+    plt.close()
+
+    # Return the saved image
+    return Image.open(save_path)
+
+# Function to load confusion matrix based on percentage and input_type
+def display_confusion_matrices_los(percentage):
+    #percentage = percentage_values_los[percentage_idx]
+    
+    # Construct folder names
+    raw_folder = os.path.join(LOS_PATH, f"raw_{percentage/100:.3f}_los_noTraining")
+    embeddings_folder = os.path.join(LOS_PATH, f"embedding_{percentage/100:.3f}_los_noTraining")
+    
+    # Process raw confusion matrix
+    raw_csv_file = os.path.join(raw_folder, f"test_predictions_raw_{percentage/100:.3f}_los.csv")
+    raw_cm_img_path = os.path.join(raw_folder, "confusion_matrix_raw.png")
+    raw_img = plot_confusion_matrix_from_csv(raw_csv_file, 
+                                             f"Confusion Matrix (Raw Channels)\n{percentage:.1f}% data", 
+                                             raw_cm_img_path)
+
+    # Process embeddings confusion matrix
+    embeddings_csv_file = os.path.join(embeddings_folder, f"test_predictions_embedding_{percentage/100:.3f}_los.csv")
+    embeddings_cm_img_path = os.path.join(embeddings_folder, "confusion_matrix_embeddings.png")
+    embeddings_img = plot_confusion_matrix_from_csv(embeddings_csv_file, 
+                                                    f"Confusion Matrix (LWM Embeddings)\n{percentage:.1f}% data", 
+                                                    embeddings_cm_img_path)
+
+    return raw_img, embeddings_img
+
+# Main function to handle user choice
+def handle_user_choice(choice, percentage=None, uploaded_file=None, emb_type='CLS Embedding'):
+    if choice == "Use Default Dataset":
+        raw_img, embeddings_img = display_confusion_matrices_los(percentage)
+        return raw_img, embeddings_img, ""  # Return empty string for console output
+    elif choice == "Upload Dataset":
+        if uploaded_file is not None:
+            raw_img, embeddings_img, console_output = process_hdf5_file(uploaded_file, percentage, emb_type)
+            return raw_img, embeddings_img, console_output
+        else:
+            return "Please upload a dataset", "Please upload a dataset", ""  # Return empty string for console output
+    else:
+        return "Invalid choice", "Invalid choice", ""  # Return empty string for console output
+
+# Custom class to capture print output
+class PrintCapture(io.StringIO):
+    def __init__(self):
+        super().__init__()
+        self.output = []
+
+    def write(self, txt):
+        self.output.append(txt)
+        super().write(txt)
+
+    def get_output(self):
+        return ''.join(self.output)
+
+# Function to load and display predefined images based on user selection
+def display_predefined_images(percentage):
+    #percentage = percentage_values_los[percentage_idx]
+    raw_image_path = os.path.join(RAW_PATH, f"percentage_{percentage}_complexity_16.png")
+    embeddings_image_path = os.path.join(EMBEDDINGS_PATH, f"percentage_{percentage}_complexity_16.png")
+    
+    # Check if the images exist
+    if os.path.exists(raw_image_path):
+        raw_image = Image.open(raw_image_path)
+    else:
+        raw_image = create_random_image()  # Use a fallback random image
+    
+    if os.path.exists(embeddings_image_path):
+        embeddings_image = Image.open(embeddings_image_path)
+    else:
+        embeddings_image = create_random_image()  # Use a fallback random image
+
+    return raw_image, embeddings_image
+
+def los_nlos_classification(file, percentage):
+    if file is not None:
+        raw_cm_image, emb_cm_image, console_output = process_hdf5_file(file, percentage)
+        return raw_cm_image, emb_cm_image, console_output  # Returning all three: two images and console output
+    else:
+        raw_image, embeddings_image = display_predefined_images(percentage)
+        return raw_image, embeddings_image, ""  # Return an empty string for console output when no file is uploaded
+
+# Function to create random images for LoS/NLoS classification results
+def create_random_image(size=(300, 300)):
+    random_image = np.random.rand(*size, 3) * 255
+    return Image.fromarray(random_image.astype('uint8'))
+
+import importlib.util
+
+# Function to dynamically load a Python module from a given file path
+def load_module_from_path(module_name, file_path):
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+# Function to split dataset into training and test sets based on user selection
+def split_dataset(channels, labels, percentage):
+    #percentage = percentage_values_los[percentage_idx] / 100
+    num_samples = channels.shape[0]
+    train_size = int(num_samples * percentage/100)
+    print(f'Number of Training Samples: {train_size}')
+    
+    indices = np.arange(num_samples)
+    np.random.shuffle(indices)
+    
+    train_idx, test_idx = indices[:train_size], indices[train_size:]
+    
+    train_data, test_data = channels[train_idx], channels[test_idx]
+    train_labels, test_labels = labels[train_idx], labels[test_idx]
+    
+    return train_data, test_data, train_labels, test_labels
+
+# Function to calculate Euclidean distance between a point and a centroid
+def euclidean_distance(x, centroid):
+    return np.linalg.norm(x - centroid)
+
+import torch
+
+def classify_based_on_distance(train_data, train_labels, test_data):
+    # Compute the centroids for the two classes
+    centroid_0 = train_data[train_labels == 0].mean(dim=0)  # Use torch.mean
+    centroid_1 = train_data[train_labels == 1].mean(dim=0)  # Use torch.mean
+    
+    predictions = []
+    for test_point in test_data:
+        # Compute Euclidean distance between the test point and each centroid
+        dist_0 = euclidean_distance(test_point, centroid_0)
+        dist_1 = euclidean_distance(test_point, centroid_1)
+        predictions.append(0 if dist_0 < dist_1 else 1)
+    
+    return torch.tensor(predictions)  # Return predictions as a PyTorch tensor
+
+def plot_confusion_matrix(y_true, y_pred, title, light_mode=False):
+    cm = confusion_matrix(y_true, y_pred)
+    
+    # Calculate F1 Score
+    f1 = f1_score(y_true, y_pred, average='weighted')
+
+    #plt.style.use('dark_background')
+    
+    # Set styling based on light or dark mode
+    if light_mode:
+        plt.style.use('default')  # Light mode styling
+        text_color = 'black'
+        cmap = 'Blues'  # Light-mode-friendly colormap
+    else:
+        plt.style.use('dark_background')  # Dark mode styling
+        text_color = 'gray'
+        #cmap = 'magma'  # Dark-mode-friendly colormap
+        cmap = 'coolwarm'
+        
+    plt.figure(figsize=(5, 5))
+    
+    # Plot the confusion matrix with a dark-mode compatible colormap
+    sns.heatmap(cm, annot=True, fmt="d", cmap=cmap, cbar=False, annot_kws={"size": 12}, linewidths=0.5, linecolor='white')
+    
+    # Add F1-score to the title
+    plt.title(f"{title}\nF1 Score: {f1:.3f}", color=text_color, fontsize=18)
+    
+    # Customize tick labels for dark mode
+    plt.xticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
+    plt.yticks([0.5, 1.5], labels=['NLoS', 'LoS'], color=text_color, fontsize=12)
+    
+    plt.ylabel('True label', color=text_color, fontsize=14)
+    plt.xlabel('Predicted label', color=text_color, fontsize=14)
+    plt.tight_layout()
+    
+    # Save the plot as an image
+    plt.savefig(f"{title}.png", transparent=True)  # Use transparent to blend with the dark mode website
+    plt.close()
+    
+    # Return the saved image
+    return Image.open(f"{title}.png")
+
+
+def identical_train_test_split(output_emb, output_raw, labels, train_percentage):
+
+    torch.manual_seed(42)
+    
+    N = output_emb.shape[0]
+    indices = torch.randperm(N) 
+    test_split_index = int(N * 0.20)
+    test_indices = indices[:test_split_index]
+    remaining_indices = indices[test_split_index:]
+    train_split_index = int(len(remaining_indices) * train_percentage / 100)
+    print(f'Training Size: {train_split_index} out of remaining {len(remaining_indices)}')
+    print(f'Test Size: {test_split_index}')
+    
+    train_indices = remaining_indices[:train_split_index]
+
+    train_emb = output_emb[train_indices]
+    test_emb = output_emb[test_indices]
+    
+    train_raw = output_raw[train_indices]
+    test_raw = output_raw[test_indices]
+
+    train_labels = labels[train_indices]
+    test_labels = labels[test_indices]
+
+    return train_emb, test_emb, train_raw, test_raw, train_labels, test_labels
+
+# Store the original working directory when the app starts
+original_dir = os.getcwd()
+
+def process_hdf5_file(uploaded_file, percentage, emb_type='CLS Embedding'):
+    capture = PrintCapture()
+    sys.stdout = capture  # Redirect print statements to capture
+    
+    try:
+        model_repo_url = "https://huggingface.co/wi-lab/lwm"
+        model_repo_dir = "./LWM"
+
+        # Step 1: Clone the repository if not already done
+        if not os.path.exists(model_repo_dir):
+            print(f"Cloning model repository from {model_repo_url}...")
+            subprocess.run(["git", "clone", model_repo_url, model_repo_dir], check=True)
+
+        # Step 2: Verify the repository was cloned and change the working directory
+        repo_work_dir = os.path.join(original_dir, model_repo_dir)
+        if os.path.exists(repo_work_dir):
+            os.chdir(repo_work_dir)  # Change the working directory only once
+            print(f"Changed working directory to {os.getcwd()}")
+            #print(f"Directory content: {os.listdir(os.getcwd())}")  # Debugging: Check repo content
+        else:
+            print(f"Directory {repo_work_dir} does not exist.")
+            return
+            
+        # Step 3: Dynamically load lwm_model.py, input_preprocess.py, and inference.py
+        lwm_model_path = os.path.join(os.getcwd(), 'lwm_model.py')
+        input_preprocess_path = os.path.join(os.getcwd(), 'input_preprocess.py')
+        inference_path = os.path.join(os.getcwd(), 'inference.py')
+
+        # Load lwm_model
+        lwm_model = load_module_from_path("lwm_model", lwm_model_path)
+
+        # Load input_preprocess
+        input_preprocess = load_module_from_path("input_preprocess", input_preprocess_path)
+
+        # Load inference
+        inference = load_module_from_path("inference", inference_path)
+
+        # Step 4: Load the model from lwm_model module
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        print(f"Loading the LWM model on {device}...")
+        model = lwm_model.lwm.from_pretrained(device=device).float()
+
+        # Step 5: Load the HDF5 file and extract the channels and labels
+        with h5py.File(uploaded_file.name, 'r') as f:
+            channels = np.array(f['channels']).astype(np.complex64)
+            labels = np.array(f['labels']).astype(np.int32)  
+        print(f"Loaded dataset with {channels.shape[0]} samples.")
+        channels = channels * (10**(-3-np.floor(np.log10(np.abs(np.mean(channels).real)))))
+        
+        # Step 7: Tokenize the data using the tokenizer from input_preprocess
+        preprocessed_chs = input_preprocess.tokenizer(manual_data=channels)
+
+        # Step 7: Perform inference using the functions from inference.py
+        if emb_type == 'Channel Embedding':
+            embedding_type = 'channel_emb'
+        elif emb_type == 'CLS Embedding':
+            embedding_type = 'cls_emb'
+            
+        output_emb = inference.lwm_inference(preprocessed_chs, embedding_type, model, device)
+        output_raw = inference.create_raw_dataset(preprocessed_chs, device)
+
+        print(f"Output Embeddings Shape: {output_emb.shape}")
+        print(f"Output Raw Shape: {output_raw.shape}")
+
+        print(f'percentage_value: {percentage}')
+        train_data_emb, test_data_emb, train_data_raw, test_data_raw, train_labels, test_labels = identical_train_test_split(output_emb.view(len(output_emb),-1),
+                                                                                                                             output_raw.view(len(output_raw),-1),
+                                                                                                                             labels,
+                                                                                                                             percentage)
+        
+        # Step 8: Perform classification using the Euclidean distance for both raw and embeddings
+        #print(f'train_data_emb: {train_data_emb.shape}')
+        #print(f'train_labels: {train_labels.shape}')
+        #print(f'test_data_emb: {test_data_emb.shape}')
+        pred_raw = classify_based_on_distance(train_data_raw, train_labels, test_data_raw)
+        pred_emb = classify_based_on_distance(train_data_emb, train_labels, test_data_emb)
+        
+        # Step 9: Generate confusion matrices for both raw and embeddings
+        raw_cm_image = plot_confusion_matrix(test_labels, pred_raw, title="Confusion Matrix (Raw Channels)")
+        emb_cm_image = plot_confusion_matrix(test_labels, pred_emb, title="Confusion Matrix (Embeddings)")
+
+        return raw_cm_image, emb_cm_image, capture.get_output()
+
+    except Exception as e:
+        return str(e), str(e), capture.get_output()
+
+    finally:
+        # Always return to the original working directory after processing
+        os.chdir(original_dir)
+        sys.stdout = sys.__stdout__  # Reset print statements
+
+######################## Define the Gradio interface ###############################
+js_code = """
+() => {
+  const isDarkMode = window.matchMedia && window.matchMedia('(prefers-color-scheme: dark)').matches;
+  return isDarkMode ? 'dark' : 'light';
+}
+"""
+
+with gr.Blocks(css="""
+    .slider-container {
+        display: inline-block;
+        margin-right: 50px;
+        text-align: center;
+    }
+
+    .explanation-box {
+        font-size: 16px;
+        font-style: italic;
+        color: #4a4a4a;
+        padding: 15px;
+        background-color: #f0f0f0;
+        border-radius: 10px;
+        margin-bottom: 20px;
+    }
+
+    .bold-highlight {
+        font-weight: bold;
+        color: #2c3e50;
+        font-size: 18px;
+        text-align: center;
+        margin-bottom: 20px;
+    }
+
+    #console-output {
+        background-color: #ffffff;  /* Light background for light mode */
+        color: #000000;  /* Dark text color for contrast */
+        padding: 10px;
+        border-radius: 5px;
+    }
+
+    .plot-title {
+        font-weight: bold;
+        color: #2c3e50;
+    }
+""") as demo:
+
+    # Contact Section
+    gr.Markdown("""
+        <div style="text-align: center;">
+            <a target="_blank" href="https://www.wi-lab.net">
+                <img src="https://www.wi-lab.net/wp-content/uploads/2021/08/WI-name.png" alt="Wireless Model" style="height: 30px;">
+            </a>
+            <a target="_blank" href="mailto:lwmwireless@gmail.com" style="margin-left: 10px;">
+                <img src="https://img.shields.io/badge/email-lwmwireless@gmail.com-blue.svg?logo=gmail" alt="Email">
+            </a>
+        </div>
+    """)
+    
+    gr.Markdown("""
+    <div class="bold-highlight">
+        🚀 Explore the pre-trained <b>LWM Model<b> here: 
+        <a target="_blank" href="https://huggingface.co/wi-lab/lwm">https://huggingface.co/wi-lab/lwm</a>
+    </div>
+    """)
+    
+    # Tab for Beam Prediction Task
+    with gr.Tab("Beam Prediction Task"):
+        #gr.Markdown("### Beam Prediction Task")
+        
+        # Explanation section with creative spacing and minimal design
+        gr.Markdown("""
+            <div style="background-color: var(--primary-background); padding: 15px; border-radius: 10px; color: var(--text-primary);">
+              <h3 style="color: var(--text-primary);">📡 <b>Sub-6GHz to mmWave Beam Prediction Task</b></h3>
+              <ul style="padding-left: 20px;">
+                <li><b>🎯 Goal</b>: Predict the strongest <b>mmWave beam</b> from a predefined codebook using Sub-6 GHz channels.</li>
+                <li><b>⚙️ Adjust Settings</b>: Use the sliders to control the training data percentage and task complexity (beam count) to explore model performance.</li>
+                <li><b>🧠 Inferences</b>:
+                  <ul>
+                    <li>🔍 First, the LWM model extracts features.</li>
+                    <li>🤖 Then, the downstream residual 1D-CNN model (500K parameters) makes beam predictions.</li>
+                  </ul>
+                </li>
+                <li><b>🗺️ Dataset</b>: A combination of six scenarios from the DeepMIMO dataset (excluded from LWM pre-training) highlights the model's strong generalization abilities.</li>
+              </ul>
+            </div>
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                data_percentage_slider = gr.Slider(label="Data Percentage for Training", minimum=10, maximum=100, step=10, value=10)
+                task_complexity_dropdown = gr.Dropdown(label="Task Complexity (Number of Beams)", choices=[16, 32, 64, 128, 256], value=16)
+                #theme_dropdown = gr.Dropdown(label="Select Theme", choices=['Light', 'Dark'], value='Light')
+
+        with gr.Row():
+            raw_img_bp = gr.Image(label="Raw Channels", type="pil", width=300, height=500)
+            embeddings_img_bp = gr.Image(label="Embeddings", type="pil", width=300, height=500)
+
+        theme_dropdown = 'Dark'
+        # Update the confusion matrices whenever sliders change
+        data_percentage_slider.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
+        task_complexity_dropdown.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
+        #theme_dropdown.change(fn=beam_prediction_task, inputs=[data_percentage_slider, task_complexity_dropdown, theme_dropdown], outputs=[raw_img_bp, embeddings_img_bp])
+        
+        # Add a conclusion section at the bottom
+        gr.Markdown("""
+            <div class="explanation-box">
+                The LWM embeddings demonstrate remarkable generalization capabilities, enabling impressive performance even with minimal training samples. This highlights their ability to effectively handle diverse tasks with limited data.
+            </div>
+        """)
+
+    # Separate Tab for LoS/NLoS Classification Task
+    with gr.Tab("LoS/NLoS Classification Task"):
+        #gr.Markdown("### LoS/NLoS Classification Task")
+
+        # Explanation section with creative spacing
+        gr.Markdown("""
+            <div style="background-color: var(--primary-background); padding: 15px; border-radius: 10px; color: var(--text-primary);">
+              <h3 style="color: var(--text-primary);">🔍 <b>LoS/NLoS Classification Task</b></h3>
+              <ul style="padding-left: 20px;">
+                <li><b>🎯 Goal</b>: Classify whether a channel is <b>LoS</b> (Line-of-Sight) or <b>NLoS</b> (Non-Line-of-Sight) with very small LWM CLS embeddings.</li>
+                <li><b>📂 Dataset</b>: Use the default dataset (a combination of six scenarios from the DeepMIMO dataset) or upload your own dataset in <b>H5</b> format.</li>
+                <li><b>💡 Custom Dataset Requirements:</b>
+                  <ul>
+                    <li>📡 <b>channels</b> array: Shape (N,32,32), rows: 32 antennas at BS, columns: 32 subcarriers</li>
+                    <li>🏷️ <b>labels</b> array: Binary LoS/NLoS values (1/0)</li>
+                  </ul>
+                </li>
+                <li><b>🔗 Tip 1</b>: Instructions for organizing your dataset are available at the bottom of the page.</li>
+                <li><b>🔗 Tip 2</b>: As the computations and inference are performed on HuggingFace CPUs, please use small datasets for faster demo experience (say <400 samples). Clone the model from <a href="https://huggingface.co/wi-lab/lwm" target="_blank">here</a> and use any number of samples locally.</li>
+                <li><b>🔗 Tip 3</b>: Your dataset will be normalized automatically based on outdoor environments. </li>
+                <li><b>💼 No Downstream Model</b>: Instead of a complex downstream model, we classify each sample based on its distance to the centroid of training samples from each class (LoS/NLoS).</il>
+              </ul>
+            </div>
+        """)
+        # Radio button for user choice: predefined data or upload dataset
+        choice_radio = gr.Radio(choices=["Use Default Dataset", "Upload Dataset"], 
+                                label="Choose how to proceed", 
+                                value="Use Default Dataset")
+        
+        percentage_slider_los = gr.Slider(minimum=float(percentage_values_los[0]), 
+                                          maximum=float(percentage_values_los[-1]), 
+                                          step=float((percentage_values_los[-1] - percentage_values_los[0]) / (len(percentage_values_los) - 1)), 
+                                          value=float(percentage_values_los[0]), 
+                                          label="Percentage of Data for Training", 
+                                          interactive=True)
+        
+        # File uploader for dataset (only visible if user chooses to upload a dataset)
+        file_input = gr.File(label="Upload HDF5 Dataset", file_types=[".h5"], visible=False)
+        
+        # Dropdown for embedding type, also only visible when "Upload Dataset" is selected
+        emb_type = gr.Dropdown(choices=["Channel Embedding", "CLS Embedding"], 
+                               value="CLS Embedding", 
+                               label="Embedding Type", 
+                               visible=False)
+        
+        # Confusion matrices display
+        with gr.Row():
+            raw_img_los = gr.Image(label="Raw Channels", type="pil", width=300, height=300)
+            embeddings_img_los = gr.Image(label="Embeddings", type="pil", width=300, height=300)
+            output_textbox = gr.Textbox(label="Console Output", lines=10, elem_id="console-output")
+        
+        # Update the visibility of file_input and emb_type based on user choice
+        def toggle_file_input_and_emb_type(choice):
+            visible = (choice == "Upload Dataset")
+            return gr.update(visible=visible), gr.update(visible=visible)
+        
+        # Change visibility of file input and embedding type dropdown based on choice
+        choice_radio.change(fn=toggle_file_input_and_emb_type, 
+                            inputs=[choice_radio], 
+                            outputs=[file_input, emb_type])
+        
+        # When percentage slider changes (for predefined data)
+        percentage_slider_los.change(fn=handle_user_choice, 
+                                     inputs=[choice_radio, percentage_slider_los, file_input, emb_type], 
+                                     outputs=[raw_img_los, embeddings_img_los, output_textbox])
+        
+        # Layout for the UI - this part does NOT need .render(), Gradio will render these automatically
+        with gr.Row():
+            file_input  # No need for .render()
+            emb_type    # No need for .render()
+
+        # Add a conclusion section at the bottom
+        gr.Markdown("""
+            <div class="explanation-box">
+                Despite their compact size (1/32 of the raw channels), LWM CLS embeddings capture rich, holistic information about the channels. This makes them exceptionally well-suited for tasks like LoS/NLoS classification, especially when working with very limited data.
+            </div>
+        """)
+
+        gr.Markdown("""
+            <div class="explanation-box">
+                To create a custom dataset, you'll need to structure your data with 32x32 channel matrices, where the rows correspond to antennas at the base station and the columns represent subcarriers. Here’s how to organize and store the channels and labels in an H5 file format for the demo:
+            </div>
+        
+            ```python
+            # How to pack channels and labels in a h5 file format as a custom dataset for the demo:
+            import h5py
+            with h5py.File('dataset.h5', 'w') as hdf:
+                hdf.create_dataset('channels', data=channels)
+                hdf.create_dataset('labels', data=labels)
+            print("Dataset saved!")
+        """)
+        gr.Markdown("""
+            <div class="explanation-box">
+                To use your preferred DeepMIMO scenarios for the custom dataset, please 
+                <a href="https://huggingface.co/wi-lab/lwm" target="_blank">clone the model and datasets</a> 
+                and follow the instructions below:
+            </div>
+
+            
+            ```python
+            from input_preprocess import DeepMIMO_data_gen deepmimo_data_cleaning label_gen # Import required modules from the model repository
+            import numpy as np
+            scenario_names = np.array([
+                "city_18_denver", "city_15_indianapolis", "city_19_oklahoma", 
+                "city_12_fortworth", "city_11_santaclara", "city_7_sandiego"
+            ])
+            scenario_name = scenario_names[0] # Select the scenario by choosing its index.
+            deepmimo_data = DeepMIMO_data_gen(scenario_name) # Generates ray-traced wireless channels for the selected scenario.
+            cleaned_deepmimo_data = deepmimo_data_cleaning(deepmimo_data) # Filters out users with no direct path to the base station (i.e., users with zero-valued channels).
+            
+            channels = np.squeeze(np.array(cleaned_deepmimo_data), axis=1) # The "channels" array is now prepared for packing into the custom dataset in H5 format.
+            labels = label_gen('LoS/NLoS Classification', deepmimo_data, scenario_name) # Generates labels for each user, classifying them as Line-of-Sight (LoS) or Non-Line-of-Sight (NLoS), and prepares the "labels" array for inclusion in the custom dataset H5 file.
+            ```
+        """)
+    #with gr.Tab("LWM Model and Framework"):    
+    #    gr.Image("images/lwm_model_v2.png")
+    #    gr.Markdown("This figure depicts the offline pre-training and online embedding generation process for LWM. The channel is divided into fixed-size patches, which are linearly embedded and combined with positional encodings before being passed through a Transformer encoder. During self-supervised pre-training, some embeddings are masked, and LWM leverages self-attention to extract deep features, allowing the decoder to reconstruct the masked values. For downstream tasks, the generated LWM embeddings enhance performance.")
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)