path change

app.py

@@ -161,26 +161,39 @@ if __name__ == "__main__":
 import gradio as gr
 import os
 import subprocess
+from pathlib import Path
+
+# Function to visualize the model
+def visualize_model(model_file):
+    # Save the uploaded file
+    file_path = os.path.join(os.getcwd(), "weight_files/yolov5.onnx") 
+    with open(file_path, "wb") as f:
+        f.write(model_file.read())
+
+    # Start Netron server to serve the file
+    port = 8081
+    subprocess.Popen(
+        ["netron", "--host", "0.0.0.0", "--port", str(port), file_path],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
 
-# Path to the model file (change this to your actual model path)
-MODEL_FILE = os.path.join(os.getcwd(), "weight_files/yolov5.onnx") 
-
-# Function to visualize the model using Netron
-def visualize_model():
-    # Serve the model using Netron
-    netron_command = f"netron.start('{MODEL_FILE}', address='0.0.0.0', port=8081)"
-    subprocess.Popen(["python3", "-c", netron_command])
-    # Return the URL where the model is served
-    return f"http://localhost:8081"
+    # Return URL for visualization
+    netron_url = f"http://localhost:{port}"
+    return f"Model uploaded successfully! View it at: {netron_url}"
 
-# Gradio Interface
-with gr.Blocks() as demo:
-    gr.Markdown("## Model Visualization with Netron")
-    output_link = gr.Textbox(label="Netron Model Link")
-    visualize_btn = gr.Button("Visualize Model")
 
-    # On button click, call the function and display the link
-    visualize_btn.click(fn=visualize_model, inputs=[], outputs=[output_link])
+# Create output directory
+os.makedirs("./models", exist_ok=True)
 
-# Launch the Gradio app
-demo.launch(server_name="0.0.0.0", server_port=7860, ssr=False)
+# Gradio Interface
+demo = gr.Interface(
+    fn=visualize_model,
+    inputs=gr.File(label="Upload your model file (ONNX, TensorFlow, etc.)"),
+    outputs=gr.Textbox(label="Netron Link"),
+    title="Model Visualizer",
+    description="Upload a model file to visualize its architecture using Netron.",
+)
+
+# Launch the app
+demo.launch(server_name="0.0.0.0", server_port=7860)