Update app.py

app.py

@@ -46,6 +46,22 @@ 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 process the uploaded .p file and perform inference using the custom model
 def process_p_file(uploaded_file, percentage_idx, complexity_idx):
     capture = PrintCapture()
@@ -64,15 +80,61 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
         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
 
-        # Simulate processing and generating random images
-        raw_image = create_random_image()
-        embeddings_image = create_random_image()
+        # 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')
+
+        print(lwm_model_path)
+        print(input_preprocess_path)
+        print(inference_path)
+        
+        # Load lwm_model
+        if os.path.exists(lwm_model_path):
+            lwm_model = load_module_from_path("lwm_model", lwm_model_path)
+        else:
+            return f"Error: lwm_model.py not found at {lwm_model_path}"
+
+        # Load input_preprocess
+        if os.path.exists(input_preprocess_path):
+            input_preprocess = load_module_from_path("input_preprocess", input_preprocess_path)
+        else:
+            return f"Error: input_preprocess.py not found at {input_preprocess_path}"
+
+        # Load inference
+        if os.path.exists(inference_path):
+            inference = load_module_from_path("inference", inference_path)
+        else:
+            return f"Error: inference.py not found at {inference_path}"
+
+        # Step 4: Load the model from lwm_model module
+        device = 'cpu'
+        print(f"Loading the LWM model on {device}...")
+        model = lwm_model.LWM.from_pretrained(device=device)
+
+        # Step 5: Tokenize the data using the tokenizer from input_preprocess
+        with open(uploaded_file.name, 'rb') as f:
+            manual_data = pickle.load(f)
+
+        preprocessed_chs = input_preprocess.tokenizer(manual_data=manual_data)
+
+        # Step 6: Perform inference using the functions from inference.py
+        output_emb = inference.lwm_inference(preprocessed_chs, 'channel_emb', model)
+        output_raw = inference.create_raw_dataset(preprocessed_chs, device)
+
+        print(f"Output Embeddings Shape: {output_emb.shape}")
+        print(f"Output Raw Shape: {output_raw.shape}")
+
+        # Step 7: Generate random images as a test
+        random_raw_image = create_random_image()
+        random_embeddings_image = create_random_image()
 
-        return raw_image, embeddings_image, capture.get_output()
+        return random_raw_image, random_embeddings_image, capture.get_output()
 
     except Exception as e:
         return str(e), str(e), capture.get_output()
@@ -85,7 +147,7 @@ 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 create_random_image(), create_random_image(), None
+        return display_predefined_images(percentage_idx, complexity_idx), None
 
 # Define the Gradio interface
 with gr.Blocks(css="""