Spaces:
Running
Running
| 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") | |
| 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] | |
| # 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 | |
| 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 | |
| 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): | |
| try: | |
| # Clone the repository if not already done (for model and tokenizer) | |
| 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) | |
| # 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 | |
| except Exception as e: | |
| return str(e), str(e) | |
| # 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) | |
| # 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:alikhani@asu.edu"><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") | |
| # Sliders for percentage and complexity | |
| 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") | |
| # 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") | |
| 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") | |
| # 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) | |
| # 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]) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() | |