Update app.py

app.py

@@ -50,72 +50,50 @@ def display_predefined_images(percentage_idx):
 
     return raw_image, embeddings_image
 
-# Updated los_nlos_classification to handle missing outputs properly
-def los_nlos_classification(file, percentage_idx):
-    if file is not None:
-        raw_cm_image, emb_cm_image, console_output = process_hdf5_file(file, percentage_idx)
-        return raw_cm_image, emb_cm_image, console_output
-    else:
-        raw_image, embeddings_image = display_predefined_images(percentage_idx)
-        return raw_image, embeddings_image, "No file uploaded. Displaying predefined images."
-
 # 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)
-    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}')
+def identical_train_test_split(output_emb, output_raw, labels, percentage_idx):
+    N = output_emb.shape[0]  # Get the total number of samples
     
-    indices = np.arange(num_samples)
-    np.random.shuffle(indices)
+    # Generate the indices for shuffling and splitting
+    indices = torch.randperm(N)  # Randomly shuffle the indices
     
-    train_idx, test_idx = indices[:train_size], indices[train_size:]
+    # 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}')
     
-    train_data, test_data = channels[train_idx], channels[test_idx]
-    train_labels, test_labels = labels[train_idx], labels[test_idx]
+    # Split indices into train and test
+    train_indices = indices[:split_index]
+    test_indices = indices[split_index:]
     
-    return train_data, test_data, train_labels, test_labels
+    # 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]
 
-# Function to calculate Euclidean distance between a point and a centroid
-def euclidean_distance(x, centroid):
-    return np.linalg.norm(x - centroid)
+    train_labels = labels[train_indices]
+    test_labels = labels[test_indices]
 
-import torch
+    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):
-    # 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
+    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:
-        # 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)
+        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)  # Return predictions as a PyTorch tensor
+    return torch.tensor(predictions)
 
 # Function to generate confusion matrix plot
 def plot_confusion_matrix(y_true, y_pred, title):
@@ -132,121 +110,75 @@ 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_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])
-    print(f'Training Size: {split_index}')
-    
-    # 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()
-
-def process_hdf5_file(uploaded_file, percentage_idx):
+# 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"
 
-        # 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 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 input_preprocess
+        # Load dynamically
+        lwm_model = load_module_from_path("lwm_model", lwm_model_path)
         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}...")
+        print(f"Loading LWM model on {device}...")
         model = lwm_model.LWM.from_pretrained(device=device).to(torch.float32)
 
-        # 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.")
-
-        # Step 7: Tokenize the data using the tokenizer from input_preprocess
+        # Preprocess and inference
         preprocessed_chs = input_preprocess.tokenizer(manual_data=channels)
-        print(preprocessed_chs[0][0][1])
-
-        # Step 7: Perform inference using the functions from inference.py
         output_emb = inference.lwm_inference(preprocessed_chs, 'channel_emb', model)
-        #print(f'output_emb:{output_emb[10][0]}')
         output_raw = inference.create_raw_dataset(preprocessed_chs, device)
-        #print(f'output_raw:{output_raw[10][0]}')
 
         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)
+        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
+        )
         
-        # 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)
-        print(f'pred_emb: {pred_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
+        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="""
@@ -262,56 +194,37 @@ with gr.Blocks(css="""
         text-align: center;
     }
 """) 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:alikhani@asu.edu" style="margin-left: 10px;">
-                <img src="https://img.shields.io/badge/email-alikhani@asu.edu-blue.svg?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=4, 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], outputs=[raw_img_bp, embeddings_img_bp])
 
+    # 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():
-            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")
-                percentage_dropdown_los = gr.Dropdown(choices=[str(v) for v in percentage_values*10], 
-                                                      value=10, 
-                                                      label="Percentage of Data for Training", 
-                                                      interactive=True)
+            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)
 
-        file_input.change(fn=los_nlos_classification, inputs=[file_input, percentage_dropdown_los], outputs=[raw_img_los, embeddings_img_los, output_textbox])
-        percentage_dropdown_los.change(fn=los_nlos_classification, inputs=[file_input, percentage_dropdown_los], outputs=[raw_img_los, embeddings_img_los, output_textbox])
+        # 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()
+    demo.launch()