app.py

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

# Paths to the predefined images folder
RAW_PATH = os.path.join("images", "raw")
EMBEDDINGS_PATH = os.path.join("images", "embeddings")

# Specific values for percentage of data for training and task complexity
percentage_values = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
complexity_values = [16, 32, 64, 128, 256]  # Task complexity values

# 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_idx, complexity_idx):
    percentage = percentage_values[percentage_idx]
    complexity = complexity_values[complexity_idx]
    
    raw_image_path = os.path.join(RAW_PATH, f"percentage_{percentage}_complexity_{complexity}.png")
    embeddings_image_path = os.path.join(EMBEDDINGS_PATH, f"percentage_{percentage}_complexity_{complexity}.png")
    
    raw_image = Image.open(raw_image_path)
    embeddings_image = Image.open(embeddings_image_path)
    
    return raw_image, embeddings_image

# 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_idx):
    percentage = percentage_values[percentage_idx] / 100
    num_samples = channels.shape[0]
    train_size = int(num_samples * percentage)
    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 classify based on distance to class centroids
def classify_based_on_distance(train_data, train_labels, test_data):
    centroid_0 = train_data[train_labels == 0].mean(dim=0)  
    centroid_1 = train_data[train_labels == 1].mean(dim=0)  
    
    predictions = []
    for test_point in test_data:
        dist_0 = torch.norm(test_point - centroid_0)
        dist_1 = torch.norm(test_point - centroid_1)
        predictions.append(0 if dist_0 < dist_1 else 1)
    
    return torch.tensor(predictions)

# Function to generate confusion matrix plot
def plot_confusion_matrix(y_true, y_pred, title):
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(5, 5))
    plt.imshow(cm, cmap='Blues')
    plt.title(title)
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.colorbar()
    plt.xticks([0, 1], labels=[0, 1])
    plt.yticks([0, 1], labels=[0, 1])
    plt.tight_layout()
    plt.savefig(f"{title}.png")
    return Image.open(f"{title}.png")

# Store the original working directory when the app starts
original_dir = os.getcwd()

# Function to process the uploaded HDF5 file for LoS/NLoS classification
def process_hdf5_file(uploaded_file, percentage_idx):
    capture = PrintCapture()
    sys.stdout = capture  # Redirect print statements to capture
    
    try:
        model_repo_url = "https://huggingface.co/sadjadalikhani/LWM"
        model_repo_dir = "./LWM"

        # Step 1: Clone the repository if not already done
        if not os.path.exists(model_repo_dir):
            subprocess.run(["git", "clone", model_repo_url, model_repo_dir], check=True)

        # Step 2: Change 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)
        else:
            print(f"Directory {repo_work_dir} does not exist.")
            return
        
        # Dynamically load the necessary modules
        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')

        lwm_model = load_module_from_path("lwm_model", lwm_model_path)
        input_preprocess = load_module_from_path("input_preprocess", input_preprocess_path)
        inference = load_module_from_path("inference", inference_path)

        device = 'cpu'
        model = lwm_model.LWM.from_pretrained(device=device)

        with h5py.File(uploaded_file.name, 'r') as f:
            channels = np.array(f['channels'])
            labels = np.array(f['labels'])

        preprocessed_chs = input_preprocess.tokenizer(manual_data=channels)

        output_emb = inference.lwm_inference(preprocessed_chs, 'channel_emb', model)
        output_raw = inference.create_raw_dataset(preprocessed_chs, device)

        return output_emb, output_raw, labels

    except Exception as e:
        return str(e), str(e), capture.get_output()

    finally:
        os.chdir(original_dir)
        sys.stdout = sys.__stdout__  

# Function to handle logic based on whether a file is uploaded or not
def los_nlos_classification(file, percentage_idx):
    if file is not None:
        return process_hdf5_file(file, percentage_idx)
    else:
        return display_predefined_images(percentage_idx), None

# Define the Gradio interface with thinner sliders
with gr.Blocks(css="""
    .slider-container {
        text-align: center;
        margin-bottom: 20px;
    }
    .image-row {
        justify-content: center;
        margin-top: 10px;
    }
    input[type=range] {
        width: 180px;
        height: 8px;
    }
""") 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:[email protected]" style="margin-left: 10px;">
                <img src="https://img.shields.io/badge/[email protected]?logo=gmail" alt="Email">
            </a>
        </div>
    """)
    
    # Tabs for Beam Prediction and LoS/NLoS Classification
    with gr.Tab("Beam Prediction Task"):
        gr.Markdown("### Beam Prediction Task")
        
        with gr.Row():
            with gr.Column(elem_id="slider-container"):
                gr.Markdown("Percentage of Data for Training")
                percentage_slider_bp = gr.Slider(minimum=0, maximum=9, step=1, value=0, label="Training Data (%)", interactive=True)
            with gr.Column(elem_id="slider-container"):
                gr.Markdown("Task Complexity")
                complexity_slider_bp = gr.Slider(minimum=0, maximum=4, step=1, value=0, label="Task Complexity", interactive=True)

        with gr.Row(elem_id="image-row"):
            raw_img_bp = gr.Image(label="Raw Channels", type="pil", width=300, height=300)
            embeddings_img_bp = gr.Image(label="Embeddings", type="pil", width=300, height=300)
        
        percentage_slider_bp.change(fn=display_predefined_images, inputs=[percentage_slider_bp, complexity_slider_bp], outputs=[raw_img_bp, embeddings_img_bp])
        complexity_slider_bp.change(fn=display_predefined_images, inputs=[percentage_slider_bp, complexity_slider_bp], outputs=[raw_img_bp, embeddings_img_bp])

    with gr.Tab("LoS/NLoS Classification Task"):
        gr.Markdown("### LoS/NLoS Classification Task")
        
        file_input = gr.File(label="Upload HDF5 Dataset", file_types=[".h5"])

        with gr.Row():
            with gr.Column(elem_id="slider-container"):
                gr.Markdown("Percentage of Data for Training")
                percentage_slider_los = gr.Slider(minimum=0, maximum=9, step=1, value=0, label="Training Data (%)", interactive=True)

        with gr.Row(elem_id="image-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=8, elem_classes="output-box")

        file_input.change(fn=los_nlos_classification, inputs=[file_input, percentage_slider_los], outputs=[raw_img_los, embeddings_img_los, output_textbox])
        percentage_slider_los.change(fn=los_nlos_classification, inputs=[file_input, percentage_slider_los], outputs=[raw_img_los, embeddings_img_los, output_textbox])

# Launch the app
if __name__ == "__main__":
    demo.launch()