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
percentage_values = np.arange(10) + 1

# 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):
    percentage = percentage_values[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

# 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'))

# Function to split dataset into training and test sets based on user selection
def identical_train_test_split(output_emb, output_raw, labels, percentage_idx):
    N = output_emb.shape[0]  # Get the total number of samples
    
    # Generate the indices for shuffling and splitting
    indices = torch.randperm(N)  # Randomly shuffle the indices
    
    # Calculate the split index
    split_index = int(N * percentage_values[percentage_idx] / 10)  # Convert percentage index to percentage value
    print(f'Training Size: {split_index}')
    
    # Split indices into train and test
    train_indices = indices[:split_index]
    test_indices = indices[split_index:]
    
    # Select the same indices from both output_emb and output_raw
    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

# Function to calculate Euclidean distance between a point and a centroid
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")

# Function to handle inference and return the results (store the results to state)
def run_inference(uploaded_file):
    capture = PrintCapture()
    sys.stdout = capture  # Redirect print statements to capture

    try:
        # Load the HDF5 file and extract channels and labels
        with h5py.File(uploaded_file.name, 'r') as f:
            channels = np.array(f['channels'])  # Assuming 'channels' dataset in the HDF5 file
            labels = np.array(f['labels'])  # Assuming 'labels' dataset in the HDF5 file
        print(f"Loaded dataset with {channels.shape[0]} samples.")

        # Run the tokenization and model inference
        model_repo_url = "https://huggingface.co/sadjadalikhani/LWM"
        model_repo_dir = "./LWM"

        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)

        # Load the model
        lwm_model_path = os.path.join(model_repo_dir, 'lwm_model.py')
        input_preprocess_path = os.path.join(model_repo_dir, 'input_preprocess.py')
        inference_path = os.path.join(model_repo_dir, 'inference.py')

        # Load dynamically
        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 = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"Loading LWM model on {device}...")
        model = lwm_model.LWM.from_pretrained(device=device).to(torch.float32)

        # Preprocess and inference
        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)

        print(f"Output Embeddings Shape: {output_emb.shape}")
        print(f"Output Raw Shape: {output_raw.shape}")

        return output_emb, output_raw, labels, capture.get_output()

    except Exception as e:
        return None, None, None, str(e)

    finally:
        sys.stdout = sys.__stdout__  # Reset print statements

# Function to handle classification after inference (using Gradio state)
def los_nlos_classification(output_emb, output_raw, labels, percentage_idx):
    if output_emb is not None and output_raw is not None:
        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_idx
        )
        
        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)

        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, "Classification successful"
    
    return create_random_image(), create_random_image(), "No valid inference outputs"

# Define the Gradio interface
with gr.Blocks(css="""
    .vertical-slider input[type=range] {
        writing-mode: bt-lr; /* IE */
        -webkit-appearance: slider-vertical; /* WebKit */
        width: 8px;
        height: 200px;
    }
    .slider-container {
        display: inline-block;
        margin-right: 50px;
        text-align: center;
    }
""") as demo:

    # Tabs for Beam Prediction and LoS/NLoS Classification
    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():
            percentage_dropdown_los = gr.Dropdown(
                choices=[str(v) for v in percentage_values * 10],
                value=10,
                label="Percentage of Data for Training",
                interactive=True
            )

        with gr.Row():
            raw_img_los = gr.Image(label="Raw Channels", type="pil", width=300, height=300, interactive=False)
            embeddings_img_los = gr.Image(label="Embeddings", type="pil", width=300, height=300, interactive=False)
            output_textbox = gr.Textbox(label="Console Output", lines=10)

        # Process file upload to run inference
        inference_output = gr.State()
        file_input.upload(run_inference, inputs=file_input, outputs=inference_output)

        # Handle dropdown change for classification
        percentage_dropdown_los.change(
            fn=los_nlos_classification,
            inputs=[inference_output['output_emb'], inference_output['output_raw'], inference_output['labels'], percentage_dropdown_los],
            outputs=[raw_img_los, embeddings_img_los, output_textbox]
        )

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