Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| from PIL import Image | |
| import numpy as np | |
| import pickle | |
| import io | |
| import sys | |
| import torch | |
| import torch | |
| import subprocess | |
| # Paths to the predefined images folder | |
| RAW_PATH = os.path.join("images", "raw") | |
| EMBEDDINGS_PATH = os.path.join("images", "embeddings") | |
| GENERATED_PATH = os.path.join("images", "generated") | |
| # Specific values for percentage and complexity | |
| 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): | |
| 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 = Image.open(raw_image_path) | |
| embeddings_image = Image.open(embeddings_image_path) | |
| return raw_image, embeddings_image | |
| # 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 | |
| # 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: | |
| model_repo_url = "https://huggingface.co/sadjadalikhani/LWM" | |
| model_repo_dir = "./LWM" | |
| 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) | |
| # Debugging: Check if the directory exists and print contents | |
| if os.path.exists(model_repo_dir): | |
| os.chdir(model_repo_dir) | |
| print(f"Changed working directory to {os.getcwd()}") | |
| print(f"Directory content: {os.listdir(os.getcwd())}") # Debugging: Check repo content | |
| else: | |
| print(f"Directory {model_repo_dir} does not exist.") | |
| return | |
| ## Step 2: Add the cloned repo to sys.path for imports | |
| #if model_repo_dir not in sys.path: | |
| # sys.path.append(model_repo_dir) | |
| ## Debugging: Print sys.path to ensure the cloned repo is in the path | |
| #print(f"sys.path: {sys.path}") | |
| sys.path.append('./') | |
| from lwm_model import LWM | |
| device = 'cpu' | |
| print(f"Loading the LWM model on {device}...") | |
| model = LWM.from_pretrained(device=device) | |
| from input_preprocess import tokenizer | |
| with open(uploaded_file.name, 'rb') as f: | |
| manual_data = pickle.load(f) | |
| preprocessed_chs = tokenizer(manual_data=manual_data) | |
| 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(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), 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: | |
| return process_p_file(file, percentage_idx, complexity_idx) | |
| else: | |
| return display_predefined_images(percentage_idx, complexity_idx), None | |
| # 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; | |
| } | |
| """) as demo: | |
| # Contact Section | |
| gr.Markdown( | |
| """ | |
| ## Contact | |
| <div style="display: flex; align-items: 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:[email protected]"><img src="https://img.shields.io/badge/[email protected]?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.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]) | |
| with gr.Tab("LoS/NLoS Classification Task"): | |
| gr.Markdown("### LoS/NLoS Classification Task") | |
| file_input = gr.File(label="Upload .p File", file_types=[".p"]) | |
| 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]) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() | |