Update app.py

app.py

@@ -159,41 +159,42 @@ if __name__ == "__main__":
     interface.launch(share=True)
 """
 import gradio as gr
-import netron
+import torch
+from ultralytics import YOLO
 import os
-import threading
-
-# 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
-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)
 
+def visualize_yolov5(model_path):
+  """
+  Visualizes the given YOLOv5 model using Netron.
+
+  Args:
+    model_path: The path to the YOLOv5 model directory.
+  """
+  try:
+    # Load the YOLOv5 model
+    model = YOLO(model_path)
+
+    # Extract the PyTorch model 
+    pytorch_model = model.model
+
+    # Save the PyTorch model to a temporary file
+    temp_model_path = "temp_model" 
+    torch.save(pytorch_model.state_dict(), os.path.join(temp_model_path, "pytorch_model.bin"))
+
+    # Run Netron 
+    os.system(f"netron {temp_model_path}") 
+
+  except Exception as e:
+    print(f"Error visualizing model: {e}")
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=visualize_yolov5,
+    inputs="text",
+    outputs="text", 
+    title="Netron YOLOv5 Model Visualization",
+    description="Enter the path to your YOLOv5 model directory."
+)
+
+# Launch the Gradio app
+iface.launch(share=True)