Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import os
 from PIL import Image
 import numpy as np
+import pickle
 
 # Paths to the predefined images folder
 RAW_PATH = os.path.join("images", "raw")
@@ -30,14 +31,7 @@ def display_predefined_images(percentage_idx, complexity_idx):
     return raw_image, embeddings_image
 
 import torch
-import numpy as np
-import importlib.util
-
-import torch
-import numpy as np
-import importlib.util
 import subprocess
-import os
 
 # Function to load the pre-trained model from your cloned repository
 def load_custom_model():
@@ -47,8 +41,8 @@ def load_custom_model():
     model.eval()  # Set the model to evaluation mode
     return model
 
-# Function to process the uploaded .py file and perform inference using the custom model
-def process_python_file(uploaded_file, percentage_idx, complexity_idx):
+# Function to process the uploaded .p file and perform inference using the custom model
+def process_p_file(uploaded_file, percentage_idx, complexity_idx):
     try:
         # Clone the repository if not already done (for model and tokenizer)
         model_repo_url = "https://huggingface.co/sadjadalikhani/LWM"
@@ -67,7 +61,6 @@ def process_python_file(uploaded_file, percentage_idx, complexity_idx):
 
         # Step 1: Load the custom model
         from lwm_model import LWM
-        import torch
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Loading the LWM model on {device}...")
         model = LWM.from_pretrained(device=device)
@@ -75,17 +68,7 @@ def process_python_file(uploaded_file, percentage_idx, complexity_idx):
         # Step 2: Import the tokenizer
         from input_preprocess import tokenizer
 
-        ## Step 3: Load the uploaded .py file that contains the wireless channel matrix
-        ## Import the Python file dynamically
-        #spec = importlib.util.spec_from_file_location("uploaded_module", uploaded_file.name)
-        #uploaded_module = importlib.util.module_from_spec(spec)
-        #spec.loader.exec_module(uploaded_module)
-
-        ## Assuming the uploaded file defines a variable called 'dataset'
-        #manual_data = uploaded_module.dataset  # This should be defined in the uploaded file
-
-        import pickle
-        # Load the .p file containing the wireless channel matrix
+        # Step 3: Load the uploaded .p file that contains the wireless channel matrix
         with open(uploaded_file.name, 'rb') as f:
             manual_data = pickle.load(f)
 
@@ -93,28 +76,11 @@ def process_python_file(uploaded_file, percentage_idx, complexity_idx):
         preprocessed_chs = tokenizer(manual_data=manual_data)
 
         # Step 5: Perform inference on the channel matrix using the model
-        #with torch.no_grad():
-        #    input_tensor = torch.tensor(preprocessed_data).unsqueeze(0)  # Add batch dimension
-        #    output = model(input_tensor)  # Perform inference
         from inference import lwm_inference, create_raw_dataset
         output_emb = lwm_inference(preprocessed_chs, 'channel_emb', model)
         output_raw = create_raw_dataset(preprocessed_chs, device)
         print(output_emb.shape)
         print(output_raw.shape)
-        # Step 6: Generate new images based on the inference results
-        #generated_raw_img = np.random.rand(300, 300, 3) * 255  # Placeholder: Replace with actual inference result
-        #generated_embeddings_img = np.random.rand(300, 300, 3) * 255  # Placeholder: Replace with actual inference result
-        
-        # Save the generated images
-        #generated_raw_image_path = os.path.join(GENERATED_PATH, f"generated_raw_{percentage_idx}_{complexity_idx}.png")
-        #generated_embeddings_image_path = os.path.join(GENERATED_PATH, f"generated_embeddings_{percentage_idx}_{complexity_idx}.png")
-        
-        #Image.fromarray(generated_raw_img.astype(np.uint8)).save(generated_raw_image_path)
-        #Image.fromarray(generated_embeddings_img.astype(np.uint8)).save(generated_embeddings_image_path)
-
-        # Load the generated images
-        #raw_image = Image.open(generated_raw_image_path)
-        #embeddings_image = Image.open(generated_embeddings_image_path)
 
         return output_emb, output_raw
 
@@ -124,8 +90,8 @@ def process_python_file(uploaded_file, percentage_idx, complexity_idx):
 # Function to handle logic based on whether a file is uploaded or not
 def los_nlos_classification(file, percentage_idx, complexity_idx):
     if file is not None:
-        # Process the uploaded file and generate new images
-        return process_python_file(file, percentage_idx, complexity_idx)
+        # Process the uploaded .p file and generate new images
+        return process_p_file(file, percentage_idx, complexity_idx)
     else:
         # Display predefined images if no file is uploaded
         return display_predefined_images(percentage_idx, complexity_idx)
@@ -150,8 +116,8 @@ with gr.Blocks(css="""
         """
         ## Contact
         <div style="display: flex; align-items: center;">
-            <a target="_blank" href="mailto:alikhani@asu.edu"><img src="https://img.shields.io/badge/email-info@wirelessmodel.com-blue.svg?logo=gmail " alt="Email"></a>&nbsp;&nbsp;
-            <a target="_blank" href="https://telegram.me/wirelessmodel"><img src="https://www.wi-lab.net/wp-content/uploads/2021/08/WI-name.png" alt="Wireless Model" style="height: 30px;"></a>&nbsp;&nbsp;
+            <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>&nbsp;&nbsp;
+            <a target="_blank" href="mailto:alikhani@asu.edu"><img src="https://img.shields.io/badge/email-alikhani@asu.edu-blue.svg?logo=gmail " alt="Email"></a>&nbsp;&nbsp;
         </div>
         """
     )
@@ -181,8 +147,8 @@ with gr.Blocks(css="""
     with gr.Tab("LoS/NLoS Classification Task"):
         gr.Markdown("### LoS/NLoS Classification Task")
         
-        # File uploader for uploading .py file
-        file_input = gr.File(label="Upload .py File", file_types=[".py"])
+        # File uploader for uploading .p file
+        file_input = gr.File(label="Upload .p File", file_types=[".p"])
 
         # Sliders for percentage and complexity
         with gr.Row():