Update app.py

app.py

@@ -160,23 +160,40 @@ if __name__ == "__main__":
 """
 import gradio as gr
 import netron
+import os
+import threading
 
-# Function to visualize the model using Netron
-def visualize_model(model_path):
-    # Start Netron to visualize the model
-    netron.start(model_path, browse=False)  # Set browse=False to avoid opening a new browser tab
-
-# Specify the local path to your model file
-model_path = os.path.join(os.getcwd(), "weight_files/yolov5.onnx")   # Change this to your actual model file path
+# Function to launch Netron in a separate thread
+def launch_netron(model_path):
+    # Start Netron server for the model file
+    netron.start(model_path)
+
+# Function to handle file upload and launch Netron
+def visualize_model(model_file):
+    # Save the uploaded model file temporarily
+    model_path = os.path.join("temp_models", model_file.name)
+    os.makedirs("temp_models", exist_ok=True)
+    
+    with open(model_path, "wb") as f:
+        f.write(model_file.read())
+    
+    # Launch Netron in a separate thread
+    threading.Thread(target=launch_netron, args=(model_path,), daemon=True).start()
+    
+    return f"Netron is running for {model_file.name}. Open your browser at http://localhost:8080."
 
 # Create Gradio interface
-iface = gr.Interface(
-    fn=lambda: visualize_model(model_path),  # Call the function without user input
-    inputs=None,  # No input required from user
-    outputs="text",  # Output type is text (or you can use "html" if needed)
-    title="Model Visualization with Netron",
-    description="This app visualizes a model using Netron."
-)
-
-# Launch the app
-iface.launch()
+with gr.Blocks() as interface:
+    gr.Markdown("<h1 style='color: lightgrey;'>Model Visualization with Netron</h1>")
+    gr.Markdown("<p style='color: lightgrey;'>Upload your model file (e.g., .onnx, .pb) to visualize it.</p>")
+    
+    model_input = gr.File(label="Upload Model File", type="file")
+    output_text = gr.Textbox(label="Output", interactive=False)
+
+    run_button = gr.Button("Visualize Model")
+    
+    run_button.click(fn=visualize_model, inputs=model_input, outputs=output_text)
+
+# Launch Gradio app
+interface.launch(share=True)
+