Update app.py

app.py

@@ -7,14 +7,16 @@ 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 and complexity
+# Specific values for percentage of data for training
 percentage_values = [10, 30, 50, 70, 100]
-complexity_values = [16, 32]
 
 # Custom class to capture print output
 class PrintCapture(io.StringIO):
@@ -30,11 +32,10 @@ class PrintCapture(io.StringIO):
         return ''.join(self.output)
 
 # Function to load and display predefined images based on user selection
-def display_predefined_images(percentage_idx, complexity_idx):
+def display_predefined_images(percentage_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_path = os.path.join(RAW_PATH, f"percentage_{percentage}_complexity_16.png")  # Assume complexity 16 for simplicity
+    embeddings_image_path = os.path.join(EMBEDDINGS_PATH, f"percentage_{percentage}_complexity_16.png")
     
     raw_image = Image.open(raw_image_path)
     embeddings_image = Image.open(embeddings_image_path)
@@ -62,8 +63,57 @@ def load_module_from_path(module_name, file_path):
     spec.loader.exec_module(module)
     return module
 
-# Function to process the uploaded .p file and perform inference using the custom model
-def process_p_file(uploaded_file, percentage_idx, complexity_idx):
+# 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 calculate Euclidean distance between a point and a centroid
+def euclidean_distance(x, centroid):
+    return np.linalg.norm(x - centroid)
+
+# Function to classify test data based on distance to class centroids
+def classify_based_on_distance(train_data, train_labels, test_data):
+    centroid_0 = np.mean(train_data[train_labels == 0], axis=0)
+    centroid_1 = np.mean(train_data[train_labels == 1], axis=0)
+    
+    predictions = []
+    for test_point in test_data:
+        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 np.array(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 process the uploaded HDF5 file and perform classification using the custom model
+def process_hdf5_file(uploaded_file, percentage_idx):
     capture = PrintCapture()
     sys.stdout = capture  # Redirect print statements to capture
     
@@ -90,51 +140,42 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
         input_preprocess_path = os.path.join(os.getcwd(), 'input_preprocess.py')
         inference_path = os.path.join(os.getcwd(), 'inference.py')
 
-        print(lwm_model_path)
-        print(input_preprocess_path)
-        print(inference_path)
-        
         # Load lwm_model
-        if os.path.exists(lwm_model_path):
-            lwm_model = load_module_from_path("lwm_model", lwm_model_path)
-        else:
-            return f"Error: lwm_model.py not found at {lwm_model_path}"
+        lwm_model = load_module_from_path("lwm_model", lwm_model_path)
 
         # Load input_preprocess
-        if os.path.exists(input_preprocess_path):
-            input_preprocess = load_module_from_path("input_preprocess", input_preprocess_path)
-        else:
-            return f"Error: input_preprocess.py not found at {input_preprocess_path}"
+        input_preprocess = load_module_from_path("input_preprocess", input_preprocess_path)
 
         # Load inference
-        if os.path.exists(inference_path):
-            inference = load_module_from_path("inference", inference_path)
-        else:
-            return f"Error: inference.py not found at {inference_path}"
+        inference = load_module_from_path("inference", inference_path)
 
         # Step 4: Load the model from lwm_model module
         device = 'cpu'
         print(f"Loading the LWM model on {device}...")
         model = lwm_model.LWM.from_pretrained(device=device)
 
-        # Step 5: Tokenize the data using the tokenizer from input_preprocess
-        with open(uploaded_file.name, 'rb') as f:
-            manual_data = pickle.load(f)
+        # 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'])  # 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.")
 
-        preprocessed_chs = input_preprocess.tokenizer(manual_data=manual_data)
+        # Step 6: Split the dataset into training and test sets
+        train_data_raw, test_data_raw, train_labels, test_labels = split_dataset(channels, labels, percentage_idx)
 
-        # Step 6: Perform inference using the functions from inference.py
-        output_emb = inference.lwm_inference(preprocessed_chs, 'channel_emb', model)
-        output_raw = inference.create_raw_dataset(preprocessed_chs, device)
+        # Step 7: Tokenize the data using the tokenizer from input_preprocess
+        preprocessed_chs = input_preprocess.tokenizer(manual_data=channels)
+        train_data_emb, test_data_emb, _, _ = split_dataset(preprocessed_chs, labels, percentage_idx)
 
-        print(f"Output Embeddings Shape: {output_emb.shape}")
-        print(f"Output Raw Shape: {output_raw.shape}")
+        # Step 8: Perform classification using the Euclidean distance for both raw and embeddings
+        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 7: Generate random images as a test
-        random_raw_image = create_random_image()
-        random_embeddings_image = create_random_image()
+        # 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 random_raw_image, random_embeddings_image, capture.get_output()
+        return raw_cm_image, emb_cm_image, capture.get_output()
 
     except Exception as e:
         return str(e), str(e), capture.get_output()
@@ -143,11 +184,11 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
         sys.stdout = sys.__stdout__  # Reset print statements
 
 # Function to handle logic based on whether a file is uploaded or not
-def los_nlos_classification(file, percentage_idx, complexity_idx):
+def los_nlos_classification(file, percentage_idx):
     if file is not None:
-        return process_p_file(file, percentage_idx, complexity_idx)
+        return process_hdf5_file(file, percentage_idx)
     else:
-        return display_predefined_images(percentage_idx, complexity_idx), None
+        return display_predefined_images(percentage_idx), None
 
 # Define the Gradio interface
 with gr.Blocks(css="""
@@ -183,38 +224,30 @@ with gr.Blocks(css="""
             with gr.Column(elem_id="slider-container"):
                 gr.Markdown("Percentage of Data for Training")
                 percentage_slider_bp = gr.Slider(minimum=0, maximum=4, step=1, value=0, interactive=True, elem_id="vertical-slider")
-            with gr.Column(elem_id="slider-container"):
-                gr.Markdown("Task Complexity")
-                complexity_slider_bp = gr.Slider(minimum=0, maximum=1, step=1, value=0, interactive=True, elem_id="vertical-slider")
 
         with gr.Row():
             raw_img_bp = gr.Image(label="Raw Channels", type="pil", width=300, height=300, interactive=False)
             embeddings_img_bp = gr.Image(label="Embeddings", type="pil", width=300, height=300, interactive=False)
 
-        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])
+        percentage_slider_bp.change(fn=display_predefined_images, inputs=[percentage_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 .p File", file_types=[".p"])
+        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=4, step=1, value=0, interactive=True, elem_id="vertical-slider")
-            with gr.Column(elem_id="slider-container"):
-                gr.Markdown("Task Complexity")
-                complexity_slider_los = gr.Slider(minimum=0, maximum=1, step=1, value=0, interactive=True, elem_id="vertical-slider")
 
         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)
 
-        file_input.change(fn=los_nlos_classification, inputs=[file_input, percentage_slider_los, complexity_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, complexity_slider_los], outputs=[raw_img_los, embeddings_img_los, output_textbox])
-        complexity_slider_los.change(fn=los_nlos_classification, inputs=[file_input, percentage_slider_los, complexity_slider_los], outputs=[raw_img_los, embeddings_img_los, output_textbox])
+        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__":