Update app.py

app.py

@@ -3,6 +3,8 @@ import os
 from PIL import Image
 import numpy as np
 import pickle
+import io
+import sys
 
 # Paths to the predefined images folder
 RAW_PATH = os.path.join("images", "raw")
@@ -13,21 +15,29 @@ GENERATED_PATH = os.path.join("images", "generated")
 percentage_values = [10, 30, 50, 70, 100]
 complexity_values = [16, 32]
 
+# 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, complexity_idx):
-    # Map the slider index to the actual value
     percentage = percentage_values[percentage_idx]
     complexity = complexity_values[complexity_idx]
-    
-    # Generate the paths to the images
     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")
     
-    # Load images using PIL
     raw_image = Image.open(raw_image_path)
     embeddings_image = Image.open(embeddings_image_path)
     
-    # Return the loaded images
     return raw_image, embeddings_image
 
 import torch
@@ -35,16 +45,17 @@ import subprocess
 
 # Function to load the pre-trained model from your cloned repository
 def load_custom_model():
-    # Assume your model is in the cloned LWM repository
     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()  # Set the model to evaluation mode
+    model.eval()
     return model
 
 # Function to process the uploaded .p file and perform inference using the custom model
 def process_p_file(uploaded_file, percentage_idx, complexity_idx):
+    capture = PrintCapture()
+    sys.stdout = capture  # Redirect print statements to capture
+
     try:
-        # Clone the repository if not already done (for model and tokenizer)
         model_repo_url = "https://huggingface.co/sadjadalikhani/LWM"
         model_repo_dir = "./LWM"
 
@@ -52,49 +63,45 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
             print(f"Cloning model repository from {model_repo_url}...")
             subprocess.run(["git", "clone", model_repo_url, model_repo_dir], check=True)
 
-        # Change the working directory to the cloned LWM folder
         if os.path.exists(model_repo_dir):
             os.chdir(model_repo_dir)
             print(f"Changed working directory to {os.getcwd()}")
         else:
             return f"Directory {model_repo_dir} does not exist."
 
-        # Step 1: Load the custom model
         from lwm_model import LWM
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Loading the LWM model on {device}...")
         model = LWM.from_pretrained(device=device)
 
-        # Step 2: Import the tokenizer
         from input_preprocess import tokenizer
 
-        # 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)
 
-        # Step 4: Tokenize the data if needed (or perform any necessary preprocessing)
         preprocessed_chs = tokenizer(manual_data=manual_data)
 
-        # Step 5: Perform inference on the channel matrix using the model
         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)
 
-        return output_emb, output_raw
+        print(f"Output Embeddings Shape: {output_emb.shape}")
+        print(f"Output Raw Shape: {output_raw.shape}")
+
+        return output_emb, output_raw, capture.get_output()
 
     except Exception as e:
-        return str(e), str(e)
+        return str(e), str(e), capture.get_output()
+
+    finally:
+        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):
     if file is not None:
-        # 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)
+        return display_predefined_images(percentage_idx, complexity_idx), None
 
 # Define the Gradio interface
 with gr.Blocks(css="""
@@ -126,7 +133,6 @@ with gr.Blocks(css="""
     with gr.Tab("Beam Prediction Task"):
         gr.Markdown("### Beam Prediction Task")
         
-        # Sliders for percentage and complexity
         with gr.Row():
             with gr.Column(elem_id="slider-container"):
                 gr.Markdown("Percentage of Data for Training")
@@ -135,22 +141,18 @@ with gr.Blocks(css="""
                 gr.Markdown("Task Complexity")
                 complexity_slider_bp = gr.Slider(minimum=0, maximum=1, step=1, value=0, interactive=True, elem_id="vertical-slider")
 
-        # Image outputs (display the images side by side and set a smaller size for the images)
         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)
 
-        # Instant image updates when sliders change
         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.Tab("LoS/NLoS Classification Task"):
         gr.Markdown("### LoS/NLoS Classification Task")
         
-        # 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():
             with gr.Column(elem_id="slider-container"):
                 gr.Markdown("Percentage of Data for Training")
@@ -159,15 +161,14 @@ with gr.Blocks(css="""
                 gr.Markdown("Task Complexity")
                 complexity_slider_los = gr.Slider(minimum=0, maximum=1, step=1, value=0, interactive=True, elem_id="vertical-slider")
 
-        # Image outputs (display the images side by side and set a smaller size for the images)
         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)
 
-        # Instant image updates based on file upload or slider changes
-        file_input.change(fn=los_nlos_classification, inputs=[file_input, percentage_slider_los, complexity_slider_los], outputs=[raw_img_los, embeddings_img_los])
-        percentage_slider_los.change(fn=los_nlos_classification, inputs=[file_input, percentage_slider_los, complexity_slider_los], outputs=[raw_img_los, embeddings_img_los])
-        complexity_slider_los.change(fn=los_nlos_classification, inputs=[file_input, percentage_slider_los, complexity_slider_los], outputs=[raw_img_los, embeddings_img_los])
+        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])
 
 # Launch the app
 if __name__ == "__main__":