Update app.py

app.py

@@ -15,9 +15,8 @@ import matplotlib.pyplot as plt
 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
+# Specific values for percentage of data for training
 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):
@@ -33,18 +32,30 @@ 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)
     
     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 load the pre-trained model from your cloned repository
+def load_custom_model():
+    from lwm_model import LWM  # Assuming the model is defined in lwm_model.py
+    model = LWM()  # Modify this according to your model initialization
+    model.eval()
+    return model
+
+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)
@@ -69,18 +80,25 @@ def split_dataset(channels, labels, percentage_idx):
     
     return train_data, test_data, train_labels, test_labels
 
-# Function to classify based on distance to class centroids
+# 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):
-    centroid_0 = train_data[train_labels == 0].mean(dim=0)  
-    centroid_1 = train_data[train_labels == 1].mean(dim=0)  
+    # 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:
-        dist_0 = torch.norm(test_point - centroid_0)
-        dist_1 = torch.norm(test_point - centroid_1)
+        # 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 torch.tensor(predictions)  # Return predictions as a PyTorch tensor
 
 # Function to generate confusion matrix plot
 def plot_confusion_matrix(y_true, y_pred, title):
@@ -97,10 +115,34 @@ def plot_confusion_matrix(y_true, y_pred, title):
     plt.savefig(f"{title}.png")
     return Image.open(f"{title}.png")
 
+def identical_train_test_split(output_emb, output_raw, labels, percentage):
+    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)
+    
+    # Split indices into train and test
+    train_indices = indices[:split_index]  # First 80% for training
+    test_indices = indices[split_index:]   # Remaining 20% for testing
+    
+    # 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
+
 # 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
@@ -111,45 +153,76 @@ def process_hdf5_file(uploaded_file, percentage_idx):
 
         # 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: Change working directory
+        # 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)
+            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
-        
-        # Dynamically load the necessary modules
+            
+        # 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 = 'cpu'
+        print(f"Loading the LWM model on {device}...")
         model = lwm_model.LWM.from_pretrained(device=device)
 
+        # 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'])
-            labels = np.array(f['labels'])
+            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.")
 
+        # 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
         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
+        print(f"Output Embeddings Shape: {output_emb.shape}")
+        print(f"Output Raw Shape: {output_raw.shape}")
+
+        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)
+        
+        # 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 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__  
+        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):
@@ -158,19 +231,18 @@ def los_nlos_classification(file, percentage_idx):
     else:
         return display_predefined_images(percentage_idx), None
 
-# Define the Gradio interface with thinner sliders
+# 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;
-        margin-bottom: 20px;
-    }
-    .image-row {
-        justify-content: center;
-        margin-top: 10px;
-    }
-    input[type=range] {
-        width: 180px;
-        height: 8px;
     }
 """) as demo:
     
@@ -193,17 +265,13 @@ with gr.Blocks(css="""
         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)
+                percentage_slider_bp = gr.Slider(minimum=0, maximum=4, step=1, value=0, interactive=True, elem_id="vertical-slider")
 
-        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.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], outputs=[raw_img_bp, embeddings_img_bp])
 
     with gr.Tab("LoS/NLoS Classification Task"):
         gr.Markdown("### LoS/NLoS Classification Task")
@@ -213,16 +281,16 @@ with gr.Blocks(css="""
         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)
+                percentage_slider_los = gr.Slider(minimum=0, maximum=4, step=1, value=0, interactive=True, elem_id="vertical-slider")
 
-        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")
+        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], 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()
+    demo.launch()